API Reference¶

Opake provides a functional API for differential privacy in PyTorch. This reference documents all public functions and classes.

Install via opake (and opake[...] extras) when using this API. Module paths remain under opake.*, but the root package is the supported user-facing installation target.

Module organization¶

Opake is organized into several modules, each focused on a specific aspect of DP training:

Core utilities¶

Serialization — Flat state_dict / from_state_dict for explicit state trees (optimizers, accounting, clip/noise state, …); template-driven restore; optional register_serializer for custom types
Exceptions — Stable semantic categories for configuration, calibration, privacy-budget, checkpoint, type, and operation failures
Random Number Generation — Immutable RNG keys for deterministic DP
RngKey — Immutable key type
key(), random_key() — Create keys
split(), fold_in() — Manipulate keys
set_reproducible_pytorch_seed() — PyTorch/cuDNN reproducibility
generator_from_key() — PyTorch generator bridge
Utilities — Functional and PyTree utilities
make_functional() — Convert nn.Module to functional form
tree_map(), tree_map_with_path(), partition(), merge(), global_norm(), tree_leaves()

DP-SGD components¶

Clipping — Per-sample gradient clipping
clipped_grad() — High-level gradient clipping
clipped_fun() — Clip and sum function outputs
clip_pytree() — Low-level PyTree clipping
adaptive_clipped_grad() — Clipped gradients with an auto-tuned clip norm
auto_clipped_grad() — AUTO-S automatic per-example gradient scaling (Bu et al. 2023)
Noise — Noise injection for DP
gaussian_noise() — Gaussian noise
mf_gaussian_noise() — Correlated noise dispatcher (DP-FTRL)
Strategy factories: band_mf_strategy(), blt_strategy(), lambda_cgd_strategy(), bisr_strategy(), identity_strategy()
Accounting — Privacy budget tracking
gaussian(), adaclip() — DP-SGD mechanisms (also via opake.dpsgd.accounting)
poisson() (plain or truncated via truncated_batch_size / dataset_size), parallel_poisson() — Poisson-family amplification
mf_gaussian(), poisson(), b_min_sep(), balls_in_bins() — DP-FTRL whole-process accounting
DpProcess operators: * (repeat), | (compose)
.epsilon_at(), .delta_at(), .advantage(), .beta_at(), .risk_at() — Privacy metrics
calibrate() — Binary-search noise multiplier for target privacy
Sampling — Privacy-amplifying sampling
PoissonSampler — Standard Poisson sampling
PoissonSampler + truncated_batch_size — Bounded Poisson sampling
CyclicPoissonSampler (opake.dpftrl) — Cyclic Poisson over bands groups; bands=1 = identity (full-data Poisson each step)
BallsInBinsSampler — Random-partition sampling (λCGD, BISR, BLT)
SequentialBatchSampler — Deterministic sequential batching (BLT)
Schedules — LR schedules for TorchOpt functional optimizers
constant_schedule() — Constant LR
cosine_schedule() — Cosine annealing
inverse_sqrt_schedule() — Inverse-square-root decay
one_minus_sqrt_schedule() — 1 - sqrt(progress) decay (concave)
with_warmup() — Compose a warmup ramp with any decay schedule
with_restarts() — Replay a schedule N times

Validation and debugging¶

Auditing — Empirical privacy validation
auditing.coin_flip(), auditing.loss_scores(), auditing.one_run() — Three-step workflow
OneRunEstimate.epsilon_at(), .delta_at(), .beta_at(), .advantage() — μ-GDP-default PLD-mirroring surface
OneRunEstimate.eps_delta() — Mechanism-agnostic (ε, δ)-DP override
attack_auc(), attack_beta_at() — Attack-side empirical metrics
Distributed — Multi-GPU training with DDP
sum_gradients() / sum_gradients_() — Copy-returning and in-place DP gradient summation
all_reduce() — Generic tensor all-reduce (sum, mean, max, min)
reduce_pytree() / reduce_pytree_() — Copy-returning and in-place generic PyTree reduction
sync() — Auto-dispatch sync for any state/aux type
local_shard() — Partition a dataset for DDP training
is_distributed(), get_rank(), get_world_size() — Distributed utilities

Hugging Face integration¶

Transformers — opake.transformers — DP-SGD for Hugging Face models
DPTrainer — full constructor, methods, callback wiring, overridable hooks
TrainingArguments — every field grouped by concern (privacy, compute, patches, save, eval, …)
opake.transformers.trainer.types — EvaluationResult, TrainOutput return types
opake.patches.apply_runtime_patches / is_runtime_patched — install/query the global runtime shims
opake.transformers.trl — TRL-style SFTTrainer / DPOTrainer (+ SFTConfig / DPOConfig), built on DPTrainer

Quick Reference¶

Typical DP-SGD Workflow¶

import opake.accounting as acc
import opake.dpsgd.accounting as dpsgd_acc
from opake.dpsgd.clipping import clipped_grad
from opake.dpsgd.noise import gaussian_noise
from opake.random import key

# Calibrate noise multiplier
result = acc.calibrate(
    acc.epsilon_budget(3.0, delta=1e-5),
    lambda nm: dpsgd_acc.poisson(dpsgd_acc.gaussian(nm), sample_rate=0.01) * 1000,
    param_min=0.1, param_max=5.0,
)

# Set up DP components
grad_fn, clip_state = clipped_grad(
    loss_fn, clipping_norm=1.0, batch_argnums=1,
    normalize_by=batch_size,
)
noise_fn, noise_state = gaussian_noise(
    noise_multiplier=result.param, key=key(42),
)

# Training loop
for batch in dataloader:
    grads, clip_state = grad_fn(params, batch, state=clip_state)
    noisy_grads, noise_state = noise_fn(grads, noise_state)
    params = update(params, noisy_grads)

See Quick Start for a complete working example.

Function index¶

Clipping¶

Function	Purpose	User Guide
`clipped_grad()`	Differentiate loss with DP clipping	Guide
`clipped_fun()`	Clip and sum function outputs	API
`clip_pytree()`	Clip PyTree to max norm	API

Noise¶

Function	Purpose	User Guide
`gaussian_noise()`	Gaussian noise	Guide
`mf_gaussian_noise()`	Correlated noise dispatcher (DP-FTRL)	Guide
`band_mf_strategy()`	BandMF banded Toeplitz strategy	Guide
`blt_strategy()`	BLT buffered Toeplitz strategy	Guide
`lambda_cgd_strategy()`	DP-λCGD PRNG-replay strategy	Guide
`bisr_strategy()`	BISR banded inverse sqrt strategy	Guide
`identity_strategy()`	Identity (DP-SGD via MF API)	Guide

Accounting (Mechanisms)¶

Function	Purpose	User Guide
`gaussian()`	Gaussian mechanism	Guide
`poisson()`	Poisson-subsampled mechanism	Guide
`poisson(..., truncated_batch_size=, dataset_size=)`	Truncated Poisson subsampling	Guide
`parallel_poisson()`	Parallel Poisson subsampling	Guide
`adaclip()`	Adaptive clipping mechanism	Guide
`eps_delta()`	Fixed (ε, δ) guarantee	Guide
`identity()`	Zero privacy loss	Guide
`nonprivate()`	Infinite-ε non-private baseline	Guide

Accounting (Matrix Factorization)¶

Function	Purpose	User Guide
`mf_gaussian()`	Strategy-driven MF Gaussian mechanism	Guide
`balls_in_bins()`	Balls-in-Bins amplification	Guide
`poisson()` (`opake.dpftrl`)	MF Poisson amplification (BandMF / identity)	Guide
`b_min_sep()`	Warm-start b-min-separation amplification	Guide

Accounting (Composition & Metrics)¶

Method / Operator	Purpose	User Guide
`process * k`	Repeat k times	Guide
`a \\| b`	Heterogeneous composition	Guide
`.epsilon_at(delta)`	Query (ε, δ)-DP	Guide
`.delta_at(epsilon)`	Query δ for given ε	Guide
`.advantage()`	Query f-DP advantage	Guide
`.beta_at(alpha)`	Query (α, β) error rates	Guide
`.risk_at(prior)`	Query Bayes risk	Guide

Calibration¶

Function	Purpose	User Guide
`calibrate()`	Find noise for target privacy	Guide
`epsilon_budget()`	Budget for (ε, δ) calibration	Guide
`delta_budget()`	Budget for δ calibration	Guide
`advantage_budget()`	Budget for advantage calibration	Guide
`beta_budget()`	Budget for (α, β) calibration	Guide
`risk_budget()`	Budget for Bayes risk calibration	Guide

Sampling¶

Class	Purpose	User Guide
`PoissonSampler`	Standard Poisson sampling	Guide
`PoissonSampler` (with `truncated_batch_size`)	Truncated Poisson sampling	Guide
`CyclicPoissonSampler` (`opake.dpftrl`)	Cyclic Poisson over `bands` groups; `bands=1` = identity	Guide
`BallsInBinsSampler`	Random-partition sampling	Guide
`SequentialBatchSampler`	Deterministic sequential batching (BLT)	Guide

Schedules¶

Function	Purpose	User Guide
`constant_schedule()`	Constant LR	Guide
`cosine_schedule()`	Cosine annealing curve	Guide
`inverse_sqrt_schedule()`	Inverse-square-root decay	Guide
`one_minus_sqrt_schedule()`	`1 - sqrt(progress)` decay (concave)	Guide
`with_warmup()`	Compose warmup ramp with a decay schedule	Guide
`with_restarts()`	Replay a schedule N times	Guide

Privacy Auditing¶

Function / Method	Purpose	User Guide
`auditing.coin_flip()`	Designate canaries + partition	Guide
`auditing.loss_scores()`	Compute membership scores	Guide
`auditing.gradient_scores()`	Compute white-box membership scores	Guide
`auditing.one_run()`	Build one-run estimate	Guide
`OneRunEstimate.epsilon_at()`	ε bound (μ-GDP default)	Guide
`OneRunEstimate.delta_at()`	δ at given ε (μ-GDP default)	Guide
`OneRunEstimate.beta_at()`	Theoretical β at α (μ-GDP)	Guide
`OneRunEstimate.advantage()`	TV at inferred μ̂ (μ-GDP)	Guide
`OneRunEstimate.eps_delta()`	(ε, δ)-DP audit method override	Guide
`OneRunEstimate.gdp(grid_size=)`	μ-GDP method with tunable grid	Guide
`OneRunEstimate.attack_auc()`	Empirical attack AUC	Guide
`OneRunEstimate.attack_beta_at()`	Empirical attack β at given FPR	Guide

Distributed¶

Function	Purpose	User Guide
`sum_gradients()` / `sum_gradients_()`	DP gradient summation (copy-returning / in-place)	Guide
`all_reduce()`	Generic tensor all-reduce (sum, mean, max, min)	Guide
`reduce_pytree()` / `reduce_pytree_()`	Generic PyTree reduction (copy-returning / in-place)	Guide
`sync()`	Auto-dispatch sync for any state/aux type	Guide
`local_shard()`	Partition dataset for DDP training	Guide
`is_distributed()`	Check if DDP is active	Guide
`get_rank()`	Get current GPU index	Guide
`get_world_size()`	Get total number of GPUs	Guide

Type hints¶

Opake uses type hints throughout. Key types:

import opake.accounting as acc

# PyTree: Nested structure of tensors
PyTree = dict[str, torch.Tensor] | tuple[torch.Tensor, ...]

# DpProcess: Composable privacy process (from Rust PLD engine)
process: DpProcess = dpsgd_acc.poisson(dpsgd_acc.gaussian(1.1), 0.01) * 1000

# Generator for reproducible noise
Generator = torch.Generator | None

Design philosophy¶

Opake's API follows these principles:

Functional-first: Immutable state, pure functions
Composable: Small, focused functions that combine naturally
Type-safe: Comprehensive type hints
PyTorch-native: Built on torch.func, works with standard PyTorch
JAX-inspired: API closely mirrors JAX-Privacy for familiarity