flax

dfd1788c - [nnx] support Array leaves in graph nodes