jax
67df04c6 - [shard-map] canonicalize rep=None to be rep={all possible axes}

Commit
297 days ago
[shard-map] canonicalize rep=None to be rep={all possible axes} None is meant to represent the same thing as {replicated over all possible axes}. But without this canonicalization, we could compare None as not equal to {all possible axes}. fixes #26621 Unrelated: in several places, including the _check_rep path, we don't handle partial auto correctly, since we treat {all possible axes} as {all mesh axes}, but actually it should be more like {all mesh axes} - auto. We'll leave that fix for a follow-up...
Author
Committer
Parents
Loading