fix: avoid NaN gradients for masked sqrt at zero

[Lyxot]

Fixes #3668

Proposed changes

Fixes NaN gradients when sqrt receives a zero cotangent at singular inputs.

For expressions like:

rho = mx.sqrt(x * x + y * y)
rho_safe = mx.maximum(rho, 1e-10)

the forward value is finite at x = y = 0, but the VJP previously evaluated the inactive sqrt(0) path as 0 / 0, producing NaN gradients.

This updates Sqrt::vjp to substitute a safe denominator when the incoming cotangent is zero, so masked-out paths contribute zero instead of NaN. The change is local to Sqrt::vjp and preserves the singular gradient when the cotangent is nonzero.

Tests

• DEVICE=cpu python -m pytest -q python/tests/test_autograd.py::TestAutograd::test_masked_sqrt_grad_is_finite
• DEVICE=gpu python -m pytest -q python/tests/test_autograd.py::TestAutograd::test_masked_sqrt_grad_is_finite
• python -m pytest -q python/tests/test_autograd.py python/tests/test_ops.py::TestOps::test_where python/tests/test_ops.py::TestOps::test_maximum python/tests/test_ops.py::TestOps::test_minimum

Checklist

Put an x in the boxes that apply.

• I have read the CONTRIBUTING document
• I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
• I have added tests that prove my fix is effective or that my feature works

[zcbenz]

Looks good to me, thanks!

[Lyxot]

Pushed a formatting commit

[angeloskath]

I am sorry to say I will close this.

First, all frameworks behave the same way. Secondly, I do think that the behavior is correct. The derivative of sqrt(x) is not really defined at 0 so if you are using sqrt(x) and x can be 0 you can't really get normal gradients.

That being said all you have to do is put the mask before the sqrt and you 're done. ie change the code to

rho_squared = x * x + y * y
rho_squared_safe = mx.maximum(rho_squared, 1e-20)
rho_safe = mx.sqrt(rho_squared_safe)