Add end-to-end example of DP training of a transformer using Flax NNX

[RamSaw]

Same as #86 but let's use Flax NNX instead of raw Jax.

[RamSaw]

debanganghosh08 do you want to take this issue and to adopt your #102 to Flax NNX? We have already an example of Flax Linen (not transformer but still), therefore Flax NNX will be more useful.

[debanganghosh08]

debanganghosh08 do you want to take this issue and to adopt your #102 to Flax NNX? We have already an example of Flax Linen (not transformer but still), therefore Flax NNX will be more useful.

Hi @RamSaw , thank you for the opportunity!
I would love to take on Issue #120. Transitioning the Transformer example to Flax NNX makes total sense as it’s the new standard, and I'm excited to help establish the pattern for using NNX with the JAX Privacy core API.

I'll get started on this right away!

[debanganghosh08]

debanganghosh08 do you want to take this issue and to adopt your #102 to Flax NNX? We have already an example of Flax Linen (not transformer but still), therefore Flax NNX will be more useful.

Hello @RamSaw ,
I agree that providing an end-to-end example in Flax NNX is the right move for the repository's long-term health. I've been researching the NNX functional API (split/merge) and believe I can provide a very clean implementation that demonstrates how to bridge NNX's stateful modules with the jax_privacy core transformations.

But I have been facing some issues as we are integrating two tools (Flax NNX and JAX Privacy) that have different philosophies. One is object oriented (stateful) and the other is functional (pure).
I am facing a ValueError: Incorrect number of leaves and the mask ndim error are both caused by a mismatch in how the model is "split" and how the data is "batched" during the gradient computation.

As per my best analysis,

1. Leaf error (42 vs 38) - When we call nnx.split(model), it separates everything into a State. If we then filter that state to just nnx.Param, we are left with 38 leaves. But the graphdef (the static part) still expects the original 42 leaves (including RNG states, and others.). To fix this, we must split the model explicitly into Params and Everything Else.
2. The Mask Error (4D vs 5D): clipped_grad uses jax.vmap internally. If my model's call expects a batch and I pass a 2D batch (B, L), vmap strips the B dimension. My model then sees a 1D input (L,). If you try to make a causal mask on a 1D input, JAX gets confused about where the batch dimension went when it tries to re-combine everything.

I will try to fix the above issues using Leaf Symmetry and Rank Consistency while using Core API and NNX Standards as you said.

Please share your opinions and any suggestions are deeply appreciated.
Thank you! 😇

[debanganghosh08]

debanganghosh08 do you want to take this issue and to adopt your #102 to Flax NNX? We have already an example of Flax Linen (not transformer but still), therefore Flax NNX will be more useful.

Hi @RamSaw ,

I wanted to provide some extra context on the architecture I chose for the NNX Transformer integration, as there were a few non-obvious JAX tracer hurdles I had to clear:

1. Trace Leakage & Dropout: Initially, I hit a TraceContextError because NNX’s reference semantics were clashing with JAX’s functional transformations. I solved this by moving the nnx.merge inside the pure_loss_fn. This ensures that layers like nnx.Dropout or nnx.MultiHeadAttention are mutating a local, transient object rather than the global model, keeping the JAX tracer happy.

2. The Masking Rank Issue: The most tricky part was that jax_privacy.clipping.clipped_grad uses vmap to compute per-example gradients. This strips the batch dimension, causing nnx.make_causal_mask to output a 3D tensor instead of the 4D tensor expected by the attention logic. I handled this by normalizing the rank within the loss function using m(x[None, :])[0], which effectively 'tricked' the attention layers into seeing a consistent batch-rank of 1.

3. Leaf Symmetry: I used the exhaustive split pattern (...) to make sure every leaf including the RNG count variables,was accounted for. This prevents those 'leaf mismatch' errors that usually crop up when the tracer re-runs the function.

✅ The code is now strictly compliant with the library's standards and should serve as a solid reference for anyone looking to use the new Flax API privately.

Please review my PR and provide an review. Thank You!