Feat : Add DP-SGD Transformer example using Flax NNX API | Issue #120 by debanganghosh08 · Pull Request #126 · google-deepmind/jax_privacy

Implemented Deplayed Preconditioners with alternating-phase protocol as per Paper 4

debanganghosh08

committed

Implement User-Level Sampling (ULS) for Transformers with per-user averaging

debanganghosh08

committed

Refactor: Use UserSelectionStrategy and clipped_grad(keep_batch_dim=False) per maintainer review

debanganghosh08

committed

style: fix remaining pylint R0917 and whitespace violations

debanganghosh08

committed

style: fix line length in user_level_transformer and sync nnx example

debanganghosh08

committed

style: fix pytype import errors and finalize production standards

debanganghosh08

committed

Refactor transformer DP-SGD example and update CI workflow

debanganghosh08

committed

Add requirements file for transformer examples to fix CI dependencies

debanganghosh08

committed

chore: align dependencies with new modular CI in pyproject.toml

debanganghosh08

committed

refactor: align Transformer examples with execution_plan and production standards

debanganghosh08

committed

fix: resolve TOML corruption and standardize formatting across all examples

debanganghosh08

committed

chore: sync infrastructure files from ULS branch

debanganghosh08

committed

refactor: revert to Adam optimizer and ensure clean TOML infrastructure

debanganghosh08

committed

chore: final production sync and docstring fix

debanganghosh08

committed

fix: restore infrastructure integrity and synchronize transformer logic

debanganghosh08

committed

Refactor: Final production hardening for NNX Transformer, secure RNG isolation, and linter compliance

debanganghosh08

committed

Fix: Align test dependencies in pyproject.toml and resolve pyink formatting for user-level examples

debanganghosh08

committed

Fix: Add Windows environment marker to flax in both test and examples dependencies

debanganghosh08

committed

Fix: Add Windows environment marker

debanganghosh08

committed

Refactor: Final production hardening, privacy-safe 2D data units, and library API alignment

debanganghosh08

committed

Fix: Resolve BandMFExecutionPlanConfig constructor error and optimize user-level buffer donation

debanganghosh08

committed

Refactor: Implement batch padding for JIT optimization, migrate to plan.clipped_grad, and patch batch_selection for multi-dim support

debanganghosh08

committed

Chore: Synchronize with upstream/main and verify JIT optimization logic

debanganghosh08

committed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Feat : Add DP-SGD Transformer example using Flax NNX API | Issue #120#126

Feat : Add DP-SGD Transformer example using Flax NNX API | Issue #120#126
debanganghosh08 wants to merge 23 commits into
google-deepmind:mainfrom
debanganghosh08:feat/nnx-transformer-dp-sgd

Commits on Apr 25, 2026

Implemented Deplayed Preconditioners with alternating-phase protocol as per Paper 4

Implement User-Level Sampling (ULS) for Transformers with per-user averaging

Refactor: Use UserSelectionStrategy and clipped_grad(keep_batch_dim=False) per maintainer review

style: fix remaining pylint R0917 and whitespace violations

style: fix line length in user_level_transformer and sync nnx example

style: fix pytype import errors and finalize production standards

Refactor transformer DP-SGD example and update CI workflow

Add requirements file for transformer examples to fix CI dependencies

chore: align dependencies with new modular CI in pyproject.toml

refactor: align Transformer examples with execution_plan and production standards

fix: resolve TOML corruption and standardize formatting across all examples

chore: sync infrastructure files from ULS branch

refactor: revert to Adam optimizer and ensure clean TOML infrastructure

chore: final production sync and docstring fix

fix: restore infrastructure integrity and synchronize transformer logic

Refactor: Final production hardening for NNX Transformer, secure RNG isolation, and linter compliance

Fix: Align test dependencies in pyproject.toml and resolve pyink formatting for user-level examples

Fix: Add Windows environment marker to flax in both test and examples dependencies

Fix: Add Windows environment marker

Refactor: Final production hardening, privacy-safe 2D data units, and library API alignment

Fix: Resolve BandMFExecutionPlanConfig constructor error and optimize user-level buffer donation

Refactor: Implement batch padding for JIT optimization, migrate to plan.clipped_grad, and patch batch_selection for multi-dim support

Chore: Synchronize with upstream/main and verify JIT optimization logic