search

google / flax

Flax is a neural network library for JAX that is designed for flexibility.
https://flax.readthedocs.io
Apache License 2.0
5.8k stars 613 forks source link

[linen] generalize transform caching #4057

Open copybara-service[bot] opened 5 days ago

copybara-service[bot] commented 5 days ago

[linen] generalize transform caching