[nnx] improve custom_vjp ergonomics

[cgarciae]

What does this PR do?

Updates the signature of the bwd function to

(*inputs, residual, output_gradient) -> inputs_tangent

Where each element in *inputs is either:

• The exact input value if it was declared as non differentiable in nondiff_argnums.
• A State object representing the gradient of state updates if the input is a graph node.
• None for all (possibly nested) JAX Arrays.

Previously output_gradient contained a (input_updates_gradient, true_output_gradient) tuple, and *inputs where only present if marked as nondiff_argnums. This change makes the output_gradient compatible with JAX while making *inputs JAX incompatible instead.

NOTE: consider following the Linen strategy of using vjp to compute the gradient wrt the desired params instead of always providing it. We would need to add nnx.vjp for this.

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB