jax

dd36a9f7 - Use `operand.aval.sharding` in convert_element_type's transpose rule