flax

Add tree-mode support for nnx.recursive_map, nnx.view, and nnx.view_info

#5267

Merged