flax
4614ac2b - support out_sharding mapping in set_metadata

Commit

2 days ago

support out_sharding mapping in set_metadata fixes #5248 PiperOrigin-RevId: 879867762

References

Author

Cristian Garcia

Committer

a-googler

Parents