jax

38b88d70 - Add the constraints in `core_call_lowering` instead of `call_lowering` which is used by shard_map too. Fixes: https://github.com/jax-ml/jax/issues/38845