Refactor BandMFExecutionPlanConfig: separate calibration from config.

The main reasoning for this change was to make it easier to calibrate a mechanism to things other than epsilon/delta.  From speaking with folks at TPDP, we probably want to support calibrating to GDP or TPR/FPR as well in the future.  This change will make that feel less clunky.

An additional benefit of this change is that the dataclass becomes significantly simpler, with epsilon/delta being removed, accountant being removed, and things like partition type and neighboring relation being automatically inferred based on the other inputs.  This should reduce cognitive load for the different classes/functions and make it easier to understand and configure a BandMF mechanism.

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()