jax.jit produces non-finite gradients for a function involving jnp.cross(x, x) and jnp.linalg.cholesky where eager gradients are finite

### Description `jax.jit` produces `nan` gradients for a function whose eager gradients are finite. The function involves `jnp.cross(x, x)` (which produces near-zero values due to float32 rounding) flowing into `jnp.linalg.cholesky`. XLA appears to miscompile the gradient computation, producing `nan` where eager correctly returns finite gradients. This is related to but distinct from the forward-pass `nan` issue with `jnp.cross(x, x)` — here the forward pass is finite in both eager and JIT, but the gradient computation diverges. ## Minimal Reproducer ```python import jax import jax.numpy as jnp _key = jax.random.PRNGKey(921742) _key, _k1 = jax.random.split(_key); t1 = jax.random.normal(_k1, (16, 3)) _key, _k2 = jax.random.split(_key); t2 = jax.random.normal(_k2, (16, 3)) _key, _k3 = jax.random.split(_key); t3 = jax.random.normal(_k3, (16, 3)) def model(t1, t2, t3): scale = t1.shape[-1] ** -0.5 w = jax.nn.softmax(jnp.matmul(t1, jnp.swapaxes(t2, -2, -1)) * scale, axis=-1) t = jnp.matmul(w, t3) c = jnp.cross(t, t) # near-zero due to float32 A = c / (jnp.linalg.norm(c) + 1.0) A = jnp.matmul(A, A.T) + jnp.eye(16) L = jnp.linalg.cholesky(A) return jnp.atleast_1d(jnp.vecdot(L, L)) grad_fn = jax.grad(lambda *a: model(*a).sum(), argnums=(0, 1, 2)) eager_grads = grad_fn(t1, t2, t3) jit_grads = jax.jit(grad_fn)(t1, t2, t3) print('forward finite:', jnp.isfinite(model(t1, t2, t3)).all()) # True print('eager grads finite:', all(jnp.isfinite(g).all() for g in eager_grads)) # True print('jit grads finite:', all(jnp.isfinite(g).all() for g in jit_grads)) # False ``` ## Expected vs Actual ``` forward finite: True eager grads finite: True ✓ correct jit grads finite: False ✗ wrong — all nan ``` ## Root Cause `jnp.cross(t, t)` with float32 produces near-zero but non-zero values (e.g. `1e-8`) due to floating point non-associativity. XLA's gradient computation through `cholesky` appears to mishandle these near-zero values, producing `nan` gradients where eager autodiff correctly returns finite values. ### System info (python version, jaxlib version, accelerator, etc.) ``` jax: 0.10.1 jaxlib: 0.10.1 numpy: 2.4.6 python: 3.12.13 | packaged by Anaconda, Inc. | (main, Mar 19 2026, 20:20:58) [GCC 14.3.0] device info: cpu-1, 1 local devices process_count: 1 platform: uname_result(system='Linux', node='unicorn-login-01', release='6.8.0-110-generic', version='#110-Ubuntu SMP PREEMPT_DYNAMIC Thu Mar 19 15:09:20 UTC 2026', machine='x86_64') JAX_PLATFORMS=cpu ```