jax
bd32431a - If the input aval mesh is empty, use the mesh of out_sharding passed to with_sharding_constraint on the output aval. We only use the mesh and not the pspec. This leads to more tracing cache hits as demonstrated in the test.

Commit

214 days ago

If the input aval mesh is empty, use the mesh of out_sharding passed to with_sharding_constraint on the output aval. We only use the mesh and not the pspec. This leads to more tracing cache hits as demonstrated in the test. PiperOrigin-RevId: 785542666

References

#30360 - If the input aval mesh is empty, use the mesh of out_sharding passed to with_sharding_constraint on the output aval. We only use the mesh and not the pspec. This leads to more tracing cache hits as demonstrated in the test.

Author

yashk2810

Committer

Google-ML-Automation

Parents

433a433a

jax bd32431a - If the input aval mesh is empty, use the mesh of out_sharding passed to with_sharding_constraint on the output aval. We only use the mesh and not the pspec. This leads to more tracing cache hits as demonstrated in the test.

jax
bd32431a - If the input aval mesh is empty, use the mesh of out_sharding passed to with_sharding_constraint on the output aval. We only use the mesh and not the pspec. This leads to more tracing cache hits as demonstrated in the test.