Add `flax.nnx.eval_shape` docstring

google / flax

Flax is a neural network library for JAX that is designed for flexibility.

https://flax.readthedocs.io

Apache License 2.0

6.15k stars 648 forks source link

Closed 8bitmp3 closed 1 week ago

8bitmp3 commented 2 weeks ago

@IvyZX @cgarciae @levskaya

version 1, can add an example + args in a separate PR.