flax
1572c84e - add nnx.map and nnx.abstract_with_sharding

Commit

141 days ago

add nnx.map and nnx.abstract_with_sharding PiperOrigin-RevId: 886959302

References

#5366 - add nnx.map and nnx.abstract_with_sharding

Author

Cristian Garcia

Committer

a-googler

a-googler

Parents

Loading