jax

f037af0c - Move check_vma check inside pvary traceable instead of it's abstract_eval rule making it a complete no-op. Before pvary as a primitive was still in the jaxpr so it's transpose rule could be triggered.

Move check_vma check inside pvary traceable instead of it's abstract_eval rule making it a complete no-op. Before pvary as a primitive was still in the jaxpr so it's transpose rule could be triggered. Also, add config._check_vma.value checks in reduced_vary_cast, preduced and vary_unreduced_cast PiperOrigin-RevId: 862911171