[shard_map] fix axis env extension bug#18138

Merged

copybara-service[bot] merged 1 commit intojax-ml:mainfrom mattjj:shmap-axis-env-fix

Oct 18, 2023

+13-1

Commits on Oct 16, 2023

[shard_map] fix axis env extension bug

mattjj
and
sharadmv
committed