jax

c7d6b653 - [sharding_in_types] Add `core.ShardingTypeError` as a new Exception that are sharding-in-types specific errors should raise.

[sharding_in_types] Add `core.ShardingTypeError` as a new Exception that are sharding-in-types specific errors should raise. This is so that we can catch this exception in backward_pass/vmap and add extra message to inform users that this is a potential JAX bug. They should file an issue on the repo. Currently, we only raise `ShardingTypeError` in one place, but we can expand to all other places in follow up changes. This change sets the machinery up. Previous error: ``` jax._src.core.ShardingTypeError: dynamic_update_slice update sharding must be equal to operand sharding, got update sharding float32[2@x]({Explicit: ('x',)}) for operand sharding float32[16]({}). ``` New error: ``` jax._src.core.ShardingTypeError: dynamic_update_slice update sharding must be equal to operand sharding, got update sharding float32[2@x]({Explicit: ('x',)}) for operand sharding float32[16]({}). This is a potential JAX bug. Please file an issue at https://github.com/jax-ml/jax/issues ``` The new added message of `This is a potential JAX bug...` is important because this error is raised in the backward pass which is 100% a JAX bug given that forward pass did not error. PiperOrigin-RevId: 739053305