flax
1f61d6b7 - Add out_sharding argument to call methods for standard layers

Commit

110 days ago

Add out_sharding argument to call methods for standard layers

References

#5102 - Add out_sharding argument to call methods for layers with jax calls that support it

Author

samanklesaria

samanklesaria

Committer

samanklesaria

samanklesaria

Parents

Loading