Skip to content
Snippets Groups Projects
Commit 969a7ef1 authored by Bart Chrzaszcz's avatar Bart Chrzaszcz Committed by TensorFlower Gardener
Browse files

#sdy Add extra `export_test.py` tests for using different meshes..

Under Shardy, we can:

- use the same mesh on save and load
- use one mesh on save and another mesh on load with different axis names
- use one mesh on save and another mesh on load with different axis names and sizes. For this case Shardy propagation may not be optimal if the module doesn't specify out shardings. This is very hard to write a unit test for, and is rare to happen, and is something we have been considering adding in Shardy b/399957785. This will be something we can allow for during Shardy propagation. This can be a standalone fix in Shardy without making any changes to JAX or XLA.

PiperOrigin-RevId: 729550047
parent ee11d428
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment