jax

a364e282 - Make _get_owner skip sharding_constraints so that the xla_metadata is applied to the op instead of the sharding constraints