jax
2ecafeee - [hijax] fix shard_map of hijax primitive

Commit

109 days ago

[hijax] fix shard_map of hijax primitive Co-authored-by: Matthew Johnson <mattjj@google.com>

References

Author

jakevdp

Parents