jax

aaac1d8f - Skip `test_spmd_preserves_input_sharding_vmap_grad` unless `xla_extension_version >= 258`