flax
866fe952 - [flax] unconditionally register nnx.Variable as a pytree

Commit

1 year ago

[flax] unconditionally register nnx.Variable as a pytree Required by changes to JAX in https://github.com/jax-ml/jax/pull/28630 PiperOrigin-RevId: 758847244

References

#4748 - [flax] unconditionally register nnx.Variable as a pytree

Author

Jake VanderPlas

Committer

a-googler

Parents

51317189

flax 866fe952 - [flax] unconditionally register nnx.Variable as a pytree

flax
866fe952 - [flax] unconditionally register nnx.Variable as a pytree