Rename differentially_private_aggregate -> optax.contrib

PiperOrigin-RevId: 628022474
google-research · Apr 25, 2024 · 3ae42de · 3ae42de
1 parent 2fb6ace
commit 3ae42de
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/scenic/train_lib/optax.py b/scenic/train_lib/optax.py
@@ -241,7 +241,7 @@ def make(config: ml_collections.ConfigDict,
       # with *zero* noise.
       grad_clip_norm_tx = [
           optax.masked(
-              optax.differentially_private_aggregate(
+              optax.contrib.differentially_private_aggregate(
                   config.max_grad_norm, 0.0, 0),
               not_frozen_mask),
           aggregate_gradients_pmean(axis_name=axis_name)]