flax

4d8c9bbe - [nnx] pytrees are graph nodes