[nnx] fix custom_vjp

[cgarciae]

What does this PR do?

• Solves the issue where FwdFn is called outside of an update context. When this happens FwdFn behaves as a pure function.
• Correctly passes the non differentiable states as captures instead of using global state.

Solves #4265.

[hrbigelow]

Hi @cgarciae - I'm excited to try this fix!

Just wanted to mention a slight discrepancy in jax vs. nnx custom_vjp signatures that are expected:

nnx bwd: res, (ins_g, outs_g) -> tangent
jax bwd: res, outs_g -> tangent

The docs here seem t obe showing an example of jax.custom_vjp.

[cgarciae]

@hrbigelow this is correct! I expanded a bit more on the docs to reflect this.

BTW: thanks for the wait. I've been trying to get this transform right but its taken some rewrites to solve some edge cases.