flax
5bcb1f56 - Use linen_vars_to_nnx_attrs in ToLinen variable restoration

Commit
10 days ago
Use linen_vars_to_nnx_attrs in ToLinen variable restoration Refactor ToLinen.__call__ to use linen_vars_to_nnx_attrs for converting Linen variables back to NNX state, replacing manual AxisMetadata unboxing and merge_state. PiperOrigin-RevId: 874740856
Author
Cristian Garcia
Committer
Parents
Loading