jax
ac87cbe0 - Remove cusolver_potrf_ffi from the list of covered backwards compatibility targets.

Commit
68 days ago
Remove cusolver_potrf_ffi from the list of covered backwards compatibility targets. This was added mistakenly in a recent change (https://github.com/jax-ml/jax/pull/32492). Once I removed it, backward compatibility tests started to fail, because the new custom call target was not covered by any test. I added tests for cholesky on GPU with the new custom call. I took the opportunity to replace the usage of pjit with jit in the export_back_compat_test. PiperOrigin-RevId: 819749498
Author
Parents
Loading