Use jax.shard_map in flax - SemanticDiff

Commit

78 days ago

Use jax.shard_map in flax PiperOrigin-RevId: 817881344

References

#5020 - Use jax.shard_map in flax

Author

yashk2810

yashk2810

Committer

a-googler

a-googler

Parents

Loading