Fix hadamard_transform on the GPU for n = m with no power-of-2 factor

[axiom-of-choice]

Fixes #4049.

Summary

hadamard_transform is documented for n = m*2^k with m in (1, 12, 20, 28), but on Metal the three sizes where the power-of-2 factor is exactly 1 (n = 12, 20, 28) failed to build the Metal library:

hadamard.h:44:19: error: constexpr variable 'num_steps' must be initialized by a constant expression
note: division by zero
  constexpr short num_steps = logN / logR;

decompose_hadamard returns (n=1, m) for those sizes, so the n2 stage is instantiated with N=1, giving logR = ctz(max_radix_2) = ctz(1) = 0 and a 0/0. The n1 and m stages are already guarded by > 1; n2 was not.

Changes

Emit and launch the n2 stage only when n2 > 1. When n2 == 1 the stage is the identity and carries no scale: scale_n2 is scale only when m == 1, and n == 1 && m == 1 returns before the primitive is created (ops.cpp:530), so n2 == 1 inside the backend always implies m > 1 and scale_n2 == 1.0. The skipped stage was also what copied x into y, so the m stage now reads x directly when n == 1; that matters on the non-donatable path, where y is freshly allocated and otherwise never written.

Cap read_width_m at n. The m kernel strides by n / read_width, so with n = 1 and read_width = 4 the group width was 1/4 == 0 and zero threads were dispatched, leaving the output untouched. Fixing only the build error surfaces this as a silent all-zero result for any batched input, so both changes are needed.

Verification

Built with Metal on an M2 Pro (Mac mini, macOS 26.5.2) and compared against the CPU backend. Bit-exact, max |gpu - cpu| = 0, for m in (12, 20, 28) across shapes (m,), (4, m), (3, 5, m), with and without an explicit scale, on both the donatable and freshly-allocated output paths. Sizes with k >= 1 (24, 40, 48, 56, 80, 112, ...) are unchanged and still bit-exact.

pytest python/tests/test_ops.py python/tests/test_autograd.py python/tests/test_vmap.py: 227 passed. The one failure, test_unstack, is unrelated and pre-existing here (np.unstack needs numpy >= 2.1, this env has 1.26.4); it fails identically without this patch. clang-format 21.1.8 and black clean.

Why it was not caught

test_hadamard builds its cases from range(1, 14), so k starts at 1 and n = m*2^0 was never tested; m = 12 and m = 20 do not appear in its [1, 28] list at all. The new test_hadamard_m_only pins all three sizes, and includes a batch dimension deliberately: the 1-D case alone still passes with the zero-thread dispatch when the input happens to be donated, so a 1-D-only test would not have caught the second half of this.

Risks

Low, and confined to n == 1. For every other size the emitted kernels, launch order, grid dimensions and read_width_m are byte-identical to before, so k >= 1 paths cannot change behavior. The one judgement call is reading from x in the m stage when n == 1; it is guarded on n > 1 ? y : x rather than changing the general path, so the multi-stage cases keep chaining through y exactly as they did.

Pre-merge / post-merge

None. No API or build changes.

[JasonHonKL]

Just personally thinks these comment are a bit nasty haha. (Maybe the author likes it idk)

[axiom-of-choice]

🤣
Can short them for sure

[JasonHonKL]

Yea! Great !

[axiom-of-choice]

Done!

[JasonHonKL]

yea just a bit too much comment to me hahaha.

[axiom-of-choice]

Done