[davrens101]

What does this PR do?

mambo_guidance (used by MagnitudeAwareGuidance) computes a per-sample norm with torch.norm(x, dim=list(range(1, ndim)), keepdim=True). With three or more dims, torch.norm is routed to linalg.matrix_norm, which only accepts a 2-tuple. So the guider fails on any normal image latent (B, C, H, W) or video latent (B, C, F, H, W) whenever guidance is active.

I found it by calling forward on every guider with random 4D tensors:

MagnitudeAwareGuidance {} 0 RUN RuntimeError linalg.matrix_norm: dim must be a 2-tuple. Got 1 2 3

This was on torch 2.15.0+cpu, on main at 33b3d37.

The fix swaps both calls for torch.linalg.vector_norm, which reduces over any set of dims and gives the same per-sample 2-norm.

I added tests/others/test_magnitude_aware_guidance.py. It compares mambo_guidance against a flatten-then-norm reference for 4D and 5D inputs in both formulations, and runs MagnitudeAwareGuidance.forward on image latents. On unmodified main, all 5 tests fail with the error above. With the change, all 5 pass. ruff check and ruff format --check are clean on the two touched files.

Limits: CPU only, on Windows. I did not run the full test suite or any end-to-end pipeline with this guider. I did not check older torch versions where torch.norm may have accepted these dims.

I used an AI coding assistant to write and test this change, and the numbers above come from its runs on my machine.