Update Keras notebooks to install jax_privacy from head

[Neal006]

Fixes #315.

Both Keras Gemma3 notebooks pinned jax_privacy==1.0.0 and imported from jax_privacy.keras import keras_api. At head the module moved to the top level (jax_privacy/keras_api.py), so that import path no longer resolves.

Changes (same in both notebooks):

• Repoint the install from jax_privacy==1.0.0 to GitHub head (pip install git+https://github.com/google-deepmind/jax_privacy.git).
• Fix the import to from jax_privacy import keras_api, matching the keras_api module docstring example.
• Drop the separate pip install dp_accounting jaxtyping drjax line: dp_accounting is now a core dependency pulled in by jax_privacy, and jaxtyping/drjax are no longer imported anywhere in the notebooks or the library.

The DPKerasConfig / make_private callsites in both notebooks still match the current API, so no config-argument changes were needed. As discussed on the issue, the flax linen MNIST notebook is intentionally pinned to 1.0 and left out of scope.

Notebooks affected:

• examples/dp_sgd_keras_gemma3_lora_finetuning_samsum.ipynb
• examples/dp_sgd_keras_gemma3_synthetic_data.ipynb

[Neal006]

Hi @ryan112358 , changes are up in #321.
One caveat on verification: I've got the code changes in, but I can't run the notebooks end to end myself. They need a TPU/GPU and Kaggle-authed Gemma3 access, and I'm on a CPU-only local setup with no Colab TPU budget to spare.
The changes are low risk (install source, import path, and dropping the now-unused dp_accounting jaxtyping drjax line), and the DPKerasConfig / make_private callsites are unchanged, so I don't expect regressions. I can do a free-tier TEST_RUN smoke pass to confirm the notebooks execute, but I can't reproduce the reported accuracy/MAUVE numbers to confirm there's no regression.
If you or someone with internal compute is better placed to do that full run, that'd be the safer path before merging. Otherwise I'll do what I can on free Colab and report back, just flagging up front that the numbers won't be a full repro.
Let me know, what are the next steps to follow in order to get the issue solved.
Thanks!

[ryan112358]

Hi @Neal006, my concern is that some of the other changes to keras_api between 1.0 and now may have introduced a regression. If you use https://colab.research.google.com/ and set runtime type to TPU, how long would it take to run the colab e2e? I think verifying it runs without error for a small number of steps would be a good spot check, but checking parity on final utility would be ideal if we can figure out how.

Adding @RamSaw to provide the formal review

[RamSaw]

Hi! Yes, this is a good change, because Colab env got updated and jax privacy 1.0.0 is not compatible anymore with the env, so these notebooks actually don't run.

The current changes LGTM, but I would like to run the updated colabs in TEST_MODE to check that it works for at least one backpropagation step, it can actually run on CPU I think (or free tier TPU of colab).

Maybe a better way to test it is to implement a CI TEST action? I was exploring it some time ago, here is how Gemini says it can look like:

"""
CI/CD Integration ("The Living Demo")

The TEST_MODE Flag: Add logic at the top of the notebook to check for an environment variable: TEST_MODE = os.environ.get("TEST_MODE") == "True". If True, truncate the dataset to ~10 samples and run only 1 epoch so it executes in minutes on a CPU.

GitHub Actions: Create a ci.yml that uses papermill to execute the notebook headlessly with the TEST_MODE=True environment variable.

Triggers: Run this CI job on two triggers:

On every PR: To ensure contributors don't merge code that breaks the demo.

Nightly Cron Job: To catch if Google Colab updates its base environment overnight and breaks our dependencies, alerting us before users see it.

Environment: Run the CI on standard ubuntu-latest (fast, free, catches 95% of logic/syntax errors), relying on !pip check to catch the Colab-specific plugin collisions.

name: Colab Demo CI

on:
  push:
    branches: [ main ]
  pull_request:
    branches: [ main ]
  schedule:
    - cron: '0 3 * * *' # Runs every night at 3 AM

jobs:
  test-notebook:
    runs-on: ubuntu-latest
    steps:
    - name: Checkout code
      uses: actions/checkout@v4

    - name: Set up Python 3.12
      uses: actions/setup-python@v5
      with:
        python-version: '3.12' # We use 3.12 because your Colab logs showed it!

    - name: Install CI tools
      run: |
        python -m pip install --upgrade pip
        pip install papermill ipykernel
        python -m ipykernel install --user --name python3

    - name: Execute Notebook in Test Mode
      env:
        TEST_MODE: "True" # We inject this as an environment variable
      run: |
        # Papermill runs demo.ipynb and saves the executed version to output.ipynb
        papermill demo.ipynb output.ipynb -k python3

"""

The other good things to do (but feel free to leave it out of scope for this PR if it is too much):

1. Check that we fix all the seeds ("Strict Seeding: Explicitly set random seeds for Python, NumPy, and JAX at the very top.")

[RamSaw]

could you add !pip check after all pip installs

same in other notebook

[RamSaw]

could you add here the note:

The results were obtained with jax_privacy=1.0.0. This version is no longer compatible with Colab environment therefore we installed jax_privacy from head in the beginning of this notebook. Due to usage of different version the final results you get might differ those documented here.

same in the other notebook

[Neal006]

Thanks @RamSaw , and thanks for confirming the compatibility break independently.
Important heads-up before anyone spends time on a TEST_MODE run: after the version bump the notebooks still don't run, and I've traced why. from jax_privacy import keras_api pulls in training.py, which fails at import with:

TypeError: jit() missing 1 required positional argument: 'fun'

Root cause is training.py:177:

@jax.jit(static_argnames=["self"], donate_argnames=["state"])

On jax 0.7.2 (what jax>=0.7.1 resolves to on current Colab), jax.jit has signature (fun, /, *, ...), so fun is required and positional-only and the bare factory form isn't valid. I reproduced it in a clean jax 0.7.2 venv: the factory form raises that exact error, and the wrapped form runs. It's the only place in the repo using the factory form; everywhere else (e.g. banded.py:241) already uses functools.partial. It came in with bff8c43.

The fix, matching the existing idiom:

@functools.partial(jax.jit, static_argnames=["self"], donate_argnames=["state"])

Since it's a library bug rather than a notebook one, I think it's cleanest as its own small PR that this one rebases on, but I'm happy to just fold it in here. Let me know which you prefer and I'll push it right away, then run both notebooks in TEST_MODE on free-tier TPU and report that they execute plus the timing Ryan asked about. Full utility parity needs the real multi-GPU run, which I can't reproduce locally, so I'd suggest keeping that as a manual check while CI guards execution.
I've also already added your two notebook asks to this PR: TEST_RUN is now driven by a TEST_MODE env var so papermill can flip it headlessly, and I added keras.utils.set_random_seed(SEED) to both notebooks (one call seeds Python, NumPy and the JAX backend).
On the CI action, I like it and I'll take it as a follow-up. Two constraints to design around: the notebooks pull Gemma3 from Kaggle, so a headless run needs KAGGLE_USERNAME/KAGGLE_KEY as repo secrets, and even the 1B TEST_MODE model is a heavy per-PR download. My suggestion is a split: a fast per-PR smoke that exercises the keras_api integration path on a tiny Keras model with no Kaggle dependency, which would have caught exactly this regression, plus an optional nightly papermill run of the full notebooks gated on those secrets. I'll open that as a separate PR so it can be reviewed on its own.

[RamSaw]

Hi! sgtm, please create a PR with a bug fix. @ryan112358 FYI that there is a bug in training.py, see above. Let's submit PR fix, run this PR, check that it works in TEST_MODE, submit this PR, then we can open a PR to add CI that runs these notebooks once per day with KAGGLE secrets (if it is possible to do so, if not, then we can give up with that, using dummy Keras model is good but not enough).