[zcbenz]

Close #4614

[ArupKantiDas]

I tested this on an M5 Pro (applegpu_g17s, so is_nax_available() is true) against a clean build of main 4ca8d5d5a. The gated delta code is the same on this PR's base 9e2e8e2 and on 4ca8d5d5a. CI never takes this path. This PR's macOS jobs ran on "Apple M2 Pro (Virtual)", and the self-hosted runners from #4795 report "Apple M4" (gen 16), so is_nax_available() is false on both.

Check (M5 Pro)	main 4ca8d5d5a	this PR ac923d1
test_fast_gated_delta.py, TF32 off (conftest)	2 failed (sequential, simdgroup)	all pass
same, MLX_ENABLE_TF32=1	2 failed	2 failed
#4614 repro, float32, dg rel. L2 vs float64 (16/48, strong decay), TF32 off	8.8e+02	2.2e-07
same, TF32 on (the default)	8.8e+02	8.8e+02
Full pytest suite	3 failed, 969 passed (the 2 above plus test_conv test_input_dilation_grad, which isn't on this PR's base yet)	0 failed, 969 passed

With TF32 off the change does what it says. float32 now takes the sequential VJP, and every input gradient in the #4614 repro is within 4e-7 of the float64 reference (16/48 and 16/16, strong and weak decay).

Some notes:

1. MLX_ENABLE_TF32 defaults to 1 (utils.h), so with default settings the #4614 repro gives the same dg error as before. bf16 inputs on the chunked VJP get it too: dg rel. L2 of 8.5e+02, against 3.2e-03 with GATED_DELTA_CHUNK_VJP=1. So the large dg error comes from the chunked dg path under strong decay, and TF32 rounding is only part of it (#4788 has a plausible explanation for why). Should #4614 stay open, or get a follow-up issue, for that part?

2. The two failures on main aren't a TF32 problem. config::get caches the first value it reads from the environment (config.cpp), and pytest runs unittest methods in alphabetical order. That means test_gated_delta_nax sets GATED_DELTA_CHUNK=16 first, and then test_gated_delta_sequential and test_gated_delta_simdgroup also run the NAX kernel, against a 1e-4 tolerance. Each passes on its own, and the file passes on main with test_gated_delta_nax deselected. On CI the cached 16 turns into 8 (no NAX), so if test_gated_delta_sequential runs after test_gated_delta_nax in the same xdist worker, it's really testing the simdgroup kernel.

3. With this PR and conftest's TF32 off, test_gated_delta_nax and the chunk=16 subtests of test_gated_delta_vjp don't reach the NAX kernels on M5 either: 16 falls back to 8 in the forward and to 1 in the VJP. The CI runners have no NAX, so right now nothing in the suite runs the NAX gated delta kernels. One way to fix both this and point 2 is to use mx.config.scoped_update(GATED_DELTA_CHUNK=...) in the three forward tests, as test_gated_delta_vjp already does, and add MLX_ENABLE_TF32=1 for the NAX cases. I tried that locally. The file passes on main and on this PR with conftest's TF32 off, and the NAX cases still hit the NAX kernels.

I used Claude Code to build both versions, run the tests and repro above.

[zcbenz]

Thanks for looking into this PR.

On the precision, at the moment we can only try to guarantee that when you use float32 and MLX_ENABLE_TF32=0 you can get float32 precision, training on float16 types is a hard problem that I don't think we would be able to look into it for a long time.

On the tests, I did some refactoring in #4783 that uses mx.config APIs to set the envs, and I have just also added MLX_ENABLE_TF32=1 to make sure NAX kernels are tested, thanks for the suggestions. I think I will merge #4783 before this PR to ensure things don't break.