Refactor BandMFExecutionPlanConfig: separate calibration from config

PiperOrigin-RevId: 926393400

Author: Ryan McKenna

examples/dp_logistic_regression.py

@@ -112,10 +112,8 @@ def main(_):
       num_bands=BANDS,
       l2_clip_norm=L2_CLIP_NORM,
       normalize_by=EXPECTED_BATCH_SIZE,
-      epsilon=EPSILON,
-      delta=DELTA,
       sampling_prob=EXPECTED_BATCH_SIZE / train_users * BANDS,
-  )
+  ).calibrate(epsilon=EPSILON, delta=DELTA)
   print('Initialized BandMFExecutionPlanConfig')
   plan = config.make()
   grad_fn = plan.clipped_grad(logistic_loss, batch_argnums=(1, 2))

examples/dp_sgd_transformer.py

@@ -247,13 +247,11 @@ def main(argv: Sequence[str]) -> None:
   config = execution_plan.BandMFExecutionPlanConfig.default(
       iterations=iterations,
       num_bands=1,
-      epsilon=epsilon,
-      delta=delta,
       sampling_prob=expected_batch_size / train_size,
       rescale_to_unit_norm=True,
       normalize_by=expected_batch_size,
       l2_clip_norm=clipping_norm,
-  )
+  ).calibrate(epsilon=epsilon, delta=delta)
   plan = config.make()
   grad_fn = plan.clipped_grad(
       loss_fn,

jax_privacy/experimental/README.md

@@ -28,8 +28,6 @@ include standard DP-SGD as the special case `num_bands=1`.
 
 ```python
 import jax.numpy as jnp
-import dp_accounting
-from jax_privacy import batch_selection
 from jax_privacy import clipping
 from jax_privacy.experimental import execution_plan
 
@@ -44,13 +42,7 @@ config = execution_plan.BandMFExecutionPlanConfig.default(
     iterations=1000,
     num_bands=1,
     sampling_prob=128 / 60000,
-    epsilon=2.0,
-    delta=1e-6,
-    partition_type=batch_selection.PartitionType.INDEPENDENT,
-    accountant=dp_accounting.pld.PLDAccountant(
-        dp_accounting.NeighboringRelation.ADD_OR_REMOVE_ONE
-    ),
-)
+).calibrate(epsilon=2.0, delta=1e-6)
 
 plan = config.make()