flax

8192b966 - [nnx] jit constrain object state