Jax code generation is not correct

I have been playing to work on Jax and NNX to generate some code examples. It is failing to generate right set of code. Any ideas how to fix this? Can I have Jax fine-tuned model that I can use on Jax code base?