flax

f63808cd - Merge pull request #3608 from google:nnx-improve-callable-proxy