Gradient Clipping¶

Per-example clipping APIs for DP training: fixed norm, AUTO-S scaling, and (on the DP-SGD side) adaptive clipping. DP-FTRL users import fixed and AUTO-S helpers from opake.dpftrl.clipping; DP-SGD users use opake.dpsgd.clipping for all three.

Overview¶

Per-sample gradient clipping bounds the influence that any single training example can have on the model, enabling differential privacy.

Clipping functions¶

clipped_grad() — High-level API: differentiates, clips, and sums gradients.
clipped_fun() — Clip and sum arbitrary function outputs (PyTrees).
auto_clipped_grad() — AUTO-S (Bu et al. NeurIPS 2023) automatic per-example gradient scaling. Constant per-record sensitivity ⇒ composes with both DP-SGD's Gaussian mechanism and DP-FTRL's matrix-factorization mechanisms.
auto_clipped_fun() — AUTO-S for arbitrary function outputs.
adaptive_clipped_grad() — Adaptive clipping (Andrew et al. 2021) with automatic threshold tuning; DP-SGD-only.
clip_pytree() — Low-level: clip an existing PyTree of gradients.
auto_scale_pytree() — Low-level: AUTO-S scale an existing PyTree.

State types¶

ClipState — Base class for clipping state.
FixedClipState — Marker state for fixed clipped_grad / clipped_fun.
AdaptiveClipState — Internal execution state for adaptive_clipped_grad.
AutoClipState — Marker state for auto_clipped_grad / auto_clipped_fun.

Auxiliary output types¶

ClippedGradAux — Per-example loss_values, grad_norms, clipped_grad_norms, loss_aux, plus aggregate clipping_rate and batch_size (from clipped_grad).
ClippedFunAux — Per-example values, norms, clipped_norms, value_aux, plus aggregate clipping_rate and batch_size (from clipped_fun).
AdaptiveClippedGradAux — Inherits all fields from ClippedGradAux (from adaptive_clipped_grad).
AutoClippedGradAux — Inherits ClippedGradAux fields with AUTO-S semantics: grad_norms (pre-scale), clipped_grad_norms (post-scale, bounded by R), clipping_rate (fraction with ‖g‖ > R), loss_values, loss_aux, group_norms, batch_size.
AutoClippedFunAux — Inherits ClippedFunAux fields with AUTO-S semantics: norms (pre-scale), clipped_norms (post-scale, bounded by R), clipping_rate (fraction with ‖v‖ > R), values, value_aux, group_norms, batch_size.

PyTree wrappers¶

ClippedPytree — Wraps a pytree of clipped gradients with calibration metadata (max_norm). Carries a max_norm which is either a scalar float or a PerGroup for per-parameter-group clipping.

noise_stddev_for(*, noise_multiplier, allocation="optimal") -> float | PerGroup Returns the noise standard deviation that gaussian_noise(noise_multiplier) would apply to this clipped pytree.

For scalar max_norm: returns noise_multiplier * max_norm.
For PerGroup max_norm with allocation="optimal" (default): returns PerGroup of per-group standard deviations σᵢ = noise_multiplier · √(Cᵢ · ΣⱼCⱼ) (MSE-optimal Mahalanobis allocation).
For PerGroup with allocation="isotropic": returns scalar noise_multiplier * max_norm.effective.

Privacy accounting is gaussian(noise_multiplier) regardless of allocation — per-group allocation costs nothing in the accountant because the Mahalanobis constraint is satisfied with equality.

grads, state = grad_fn(params, x, state=clip_state)
stddev = grads.noise_stddev_for(noise_multiplier=0.8)
noise_fn, noise_state = gaussian_noise(noise_multiplier=0.8, ...)

NoisedPytree — Subclass of ClippedPytree that also carries a noise_stddev field (scalar or PerGroup). Returned by all noise functions.

Distributed sync helpers¶

Use sync() from opake.distributed to synchronize any clipping state or aux object. It auto-dispatches to the right function based on type:

sync(FixedClipState | AutoClipState) → marker-state passthrough.
sync(AdaptiveClipState) → sums local clipped counts and batch sizes across ranks, uses the centered statistic with the public expected batch size, and updates the internal next adaptive threshold identically on every rank.
sync(ClippedFunAux | ClippedGradAux | AutoClippedFunAux | AutoClippedGradAux | AdaptiveClippedGradAux) → gathers aux across ranks.

Dispatch resolves by exact type and then by __mro__, so any subclass of a registered clipping state or aux type reaches its base handler; an unregistered type raises InputTypeError rather than passing through unsynchronized.

API Documentation¶

`clipping` ¶

DP-SGD clipping façade — fixed, AUTO-S, per-group, and adaptive thresholding.

Headline factories:

:func:clipped_grad — fixed-threshold per-example clipping.
:func:auto_clipped_grad — AUTO-S (Bu et al., Automatic Clipping <https://arxiv.org/abs/2206.07136>_, NeurIPS 2023).
:func:per_group — build :class:~opake.types.PerGroup groupings.
:func:adaptive_clipped_grad — adaptive clipping (Andrew et al., Differentially Private Learning with Adaptive Clipping <https://arxiv.org/abs/1905.03871>_, 2021).

Adaptive clipping is DP-SGD-only: its threshold drifts across steps, which violates the constant-sensitivity assumption matrix-factorization proofs require.

State and auxiliary dataclasses live in :mod:opake.dpsgd.clipping.types. AUTO-S function-level helpers live in :mod:opake.dpsgd.clipping.fun.

`adaptive_clipped_grad(loss_fn, argnums=0, has_aux=False, *, initial_clipping_norm=0.1, target_quantile=0.5, learning_rate=0.2, clipping_norm_min=0.01, clipping_norm_max=100.0, fraction_noise_std=_DEFAULT_FRACTION_NOISE_STD, expected_batch_size=None, key, return_aux=False, return_stats=False, pre_clipping_transform=lambda x: x, **clipped_grad_kwargs)` ¶

Create function for adaptive gradient clipping with explicit state-passing.

Returns (clipped_grad_fn, initial_state). The returned function takes state as an explicit parameter and returns (grad, new_state) or ((grad, aux), new_state).

Parameters:

Name	Type	Description	Default
`loss_fn`	`Callable`	Loss function (scalar output). If `has_aux`, returns `(scalar, loss_aux)`.	required
`argnums`	`int \| tuple[int, ...]`	Which argument(s) to differentiate w.r.t.	`0`
`has_aux`	`bool`	If True, `loss_fn` returns `(value, loss_aux)`.	`False`
`initial_clipping_norm`	`float \| PerGroup`	Initial clipping threshold C_0. When `PerGroup`, each group adapts independently.	`0.1`
`target_quantile`	`float`	Target fraction of clipped gradients.	`0.5`
`learning_rate`	`float`	Step size for geometric adaptation.	`0.2`
`clipping_norm_min`	`float`	Minimum allowed clipping threshold.	`0.01`
`clipping_norm_max`	`float`	Maximum allowed clipping threshold.	`100.0`
`fraction_noise_std`	`float`	Std of Gaussian noise added to the centered clipping indicator (default 0.05).	`_DEFAULT_FRACTION_NOISE_STD`
`expected_batch_size`	`float \| None`	Public, fixed denominator for the clipping indicator. Defaults to an explicitly supplied `normalize_by`. Pass the same value to the AdaClip accountant.	`None`
`key`	`RngKey`	RNG key for quantile noise generation.	required
`return_aux`	`bool`	If True, return per-example aux with loss values, gradient norms, and clipping rate.	`False`
`return_stats`	`bool`	If True, return aggregate clipping statistics alongside the gradient. Cannot be combined with `return_aux`.	`False`
`pre_clipping_transform`	`Callable`	An optional function to apply to the per-example gradients before clipping. The function should consume the gradient pytree for a single example and return a new pytree (possibly with different structure). Can be used to e.g., scale the leaves of the pytree to accommodate preconditioner clipping. Does not affect the sensitivity guarantee. Default is identity function.	`lambda x: x`
`**clipped_grad_kwargs`	`Any`	Passed to `clipped_grad()` (`batch_argnums`, `normalize_by`, etc).	`{}`

Note

Empty-batch parity (second_moment): the empty-batch short-circuit mirrors :func:~opake.dpsgd.clipping.clipped_grad and :func:~opake.dpsgd.clipping.auto_clipped_grad — when second_moment=True it returns a :class:~opake.types.SecondMomentClippingOutput of zeros (squared-stream sensitivity C²/normalize_by), so paired-stream noise and optimizer dispatch are stable across empty and non-empty steps under Poisson sampling.

Returns:

Type	Description
`tuple[Callable, AdaptiveClipState]`	A tuple of (clipped_grad_fn, initial_state) where: - clipped_grad_fn: Function with signature (args, state, kwargs) -> (grad, new_state) or (args, state, **kwargs) -> ((grad, aux), new_state) or `((grad, stats), new_state)` when `return_stats=True`. - initial_state: Initial AdaptiveClipState

Example (single device): >>> import torch >>> from opake.dpsgd.clipping import adaptive_clipped_grad >>> from opake.dpsgd.noise import gaussian_noise >>> from opake.random import key >>> import torchopt >>> >>> def loss_fn(params, x, y): ... pred = x @ params ... return ((pred - y) ** 2).mean() >>> >>> # Create an adaptive clipping function with explicit local state. >>> # Distributed callers must synchronize state explicitly. >>> grad_fn, clip_state = adaptive_clipped_grad( ... loss_fn, ... initial_clipping_norm=0.1, ... target_quantile=0.5, ... key=key(0), ... batch_argnums=(1, 2), ... normalize_by=4, ... ) >>> >>> # Training loop with explicit state-passing >>> params = torch.randn(10, requires_grad=False) >>> optimizer = torchopt.adamw(lr=1e-3) >>> opt_state = optimizer.init(params) >>> >>> noise_fn, noise_state = gaussian_noise(noise_multiplier=1.1, key=key(1)) >>> # dataloader yields (x, y) batches (e.g. a torch DataLoader). >>> dataloader = [(torch.randn(4, 10), torch.randn(4)) for _ in range(3)] >>> for batch_x, batch_y in dataloader: ... # Compute clipped gradients - state passed explicitly ... grad, clip_state = grad_fn(params, batch_x, batch_y, state=clip_state) ... ... # Add DP noise scaled to current threshold ... noisy_grad, noise_state = noise_fn(grad, noise_state) ... ... # Optimizer step ... updates, opt_state = optimizer.update(noisy_grad, opt_state, params=params) ... params = torchopt.apply_updates(params, updates) ... ... # Monitor adaptation ... # The current DP max_norm is attached to the clipped output. ... current_bound = grad.max_norm

Example with distributed training (DDP with Poisson sampling): >>> import torch.distributed as dist >>> from torch.utils.data import DataLoader >>> from opake.dpsgd.clipping import adaptive_clipped_grad >>> from opake.dpsgd.noise import gaussian_noise >>> from opake.dpsgd.sampling import PoissonSampler >>> from opake.distributed import local_shard, sum_gradients, sync >>> from opake.random import fold_in, key >>> >>> # Initialize distributed first (typically via torchrun). >>> dist.init_process_group(backend='nccl') >>> rank, world_size = dist.get_rank(), dist.get_world_size() >>> >>> # Create adaptive clipping with the global public expected batch size. >>> expected_batch_size = 0.01 * len(dataset) >>> grad_fn, clip_state = adaptive_clipped_grad( ... loss_fn, ... key=key(0), ... batch_argnums=(1, 2), ... normalize_by=expected_batch_size, ... ) >>> noise_fn, noise_state = gaussian_noise(noise_multiplier=1.1, key=key(2)) >>> >>> # Poisson sampling: shard per rank, derive a per-rank sampling key. >>> shard = local_shard(dataset, rank=rank, world_size=world_size) >>> sampler = PoissonSampler( ... shard, sample_rate=0.01, key=fold_in(key(42), rank) ... ) >>> loader = DataLoader(shard, batch_sampler=sampler) >>> >>> for batch_x, batch_y in loader: ... # Each device: compute clipped gradients and local adaptive state ... grad, clip_state = grad_fn(params, batch_x, batch_y, state=clip_state) ... # Explicit sync step for adaptive clipping state. ... clip_state = sync(clip_state) ... ... # Sum clipped gradients across devices ... grad = sum_gradients(grad) ... ... # Add noise and update ... noisy_grad, noise_state = noise_fn(grad, noise_state) ... # ... optimizer step

References

Andrew et al., "Differentially Private Learning with Adaptive Clipping", NeurIPS 2021. https://arxiv.org/abs/1905.03871

Source code in packages/opake-dpsgd/src/opake/api/dpsgd/clipping/_adaptive.py

def adaptive_clipped_grad(
    loss_fn: Callable,
    argnums: int | tuple[int, ...] = 0,
    has_aux: bool = False,
    *,
    initial_clipping_norm: float | PerGroup = 0.1,
    target_quantile: float = 0.5,
    learning_rate: float = 0.2,
    clipping_norm_min: float = 0.01,
    clipping_norm_max: float = 100.0,
    fraction_noise_std: float = _DEFAULT_FRACTION_NOISE_STD,
    expected_batch_size: float | None = None,
    key: RngKey,
    return_aux: bool = False,
    return_stats: bool = False,
    pre_clipping_transform: Callable = lambda x: x,
    **clipped_grad_kwargs: Any,
) -> tuple[Callable, AdaptiveClipState]:
    """Create function for adaptive gradient clipping with explicit state-passing.

    Returns ``(clipped_grad_fn, initial_state)``.  The returned function
    takes ``state`` as an explicit parameter and returns
    ``(grad, new_state)`` or ``((grad, aux), new_state)``.

    Args:
        loss_fn: Loss function (scalar output). If ``has_aux``, returns
            ``(scalar, loss_aux)``.
        argnums: Which argument(s) to differentiate w.r.t.
        has_aux: If True, ``loss_fn`` returns ``(value, loss_aux)``.
        initial_clipping_norm: Initial clipping threshold C_0.  When
            ``PerGroup``, each group adapts independently.
        target_quantile: Target fraction of clipped gradients.
        learning_rate: Step size for geometric adaptation.
        clipping_norm_min: Minimum allowed clipping threshold.
        clipping_norm_max: Maximum allowed clipping threshold.
        fraction_noise_std: Std of Gaussian noise added to the centered
            clipping indicator (default 0.05).
        expected_batch_size: Public, fixed denominator for the clipping
            indicator. Defaults to an explicitly supplied ``normalize_by``.
            Pass the same value to the AdaClip accountant.
        key: RNG key for quantile noise generation.
        return_aux: If True, return per-example aux with loss values,
            gradient norms, and clipping rate.
        return_stats: If True, return aggregate clipping statistics alongside
            the gradient. Cannot be combined with ``return_aux``.
        pre_clipping_transform: An optional function to apply to the
            per-example gradients before clipping. The function should
            consume the gradient pytree for a single example and return
            a new pytree (possibly with different structure). Can be
            used to e.g., scale the leaves of the pytree to accommodate
            preconditioner clipping. Does not affect the sensitivity guarantee.
            Default is identity function.
        **clipped_grad_kwargs: Passed to ``clipped_grad()``
            (``batch_argnums``, ``normalize_by``, etc).

    Note:
        **Empty-batch parity (``second_moment``):** the empty-batch
        short-circuit mirrors :func:`~opake.dpsgd.clipping.clipped_grad` and
        :func:`~opake.dpsgd.clipping.auto_clipped_grad` — when ``second_moment=True``
        it returns a :class:`~opake.types.SecondMomentClippingOutput` of zeros
        (squared-stream sensitivity ``C²/normalize_by``), so paired-stream
        noise and optimizer dispatch are stable across empty and non-empty
        steps under Poisson sampling.

    Returns:
        A tuple of (clipped_grad_fn, initial_state) where:
            - clipped_grad_fn: Function with signature
                (*args, state, **kwargs) -> (grad, new_state) or
                (*args, state, **kwargs) -> ((grad, aux), new_state)
                or ``((grad, stats), new_state)`` when ``return_stats=True``.
            - initial_state: Initial AdaptiveClipState

    Example (single device):
        >>> import torch
        >>> from opake.dpsgd.clipping import adaptive_clipped_grad
        >>> from opake.dpsgd.noise import gaussian_noise
        >>> from opake.random import key
        >>> import torchopt
        >>>
        >>> def loss_fn(params, x, y):
        ...     pred = x @ params
        ...     return ((pred - y) ** 2).mean()
        >>>
        >>> # Create an adaptive clipping function with explicit local state.
        >>> # Distributed callers must synchronize state explicitly.
        >>> grad_fn, clip_state = adaptive_clipped_grad(
        ...     loss_fn,
        ...     initial_clipping_norm=0.1,
        ...     target_quantile=0.5,
        ...     key=key(0),
        ...     batch_argnums=(1, 2),
        ...     normalize_by=4,
        ... )
        >>>
        >>> # Training loop with explicit state-passing
        >>> params = torch.randn(10, requires_grad=False)
        >>> optimizer = torchopt.adamw(lr=1e-3)
        >>> opt_state = optimizer.init(params)
        >>>
        >>> noise_fn, noise_state = gaussian_noise(noise_multiplier=1.1, key=key(1))
        >>> # ``dataloader`` yields (x, y) batches (e.g. a torch DataLoader).
        >>> dataloader = [(torch.randn(4, 10), torch.randn(4)) for _ in range(3)]
        >>> for batch_x, batch_y in dataloader:
        ...     # Compute clipped gradients - state passed explicitly
        ...     grad, clip_state = grad_fn(params, batch_x, batch_y, state=clip_state)
        ...
        ...     # Add DP noise scaled to current threshold
        ...     noisy_grad, noise_state = noise_fn(grad, noise_state)
        ...
        ...     # Optimizer step
        ...     updates, opt_state = optimizer.update(noisy_grad, opt_state, params=params)
        ...     params = torchopt.apply_updates(params, updates)
        ...
        ...     # Monitor adaptation
        ...     # The current DP max_norm is attached to the clipped output.
        ...     current_bound = grad.max_norm

    Example with distributed training (DDP with Poisson sampling):
        >>> import torch.distributed as dist
        >>> from torch.utils.data import DataLoader
        >>> from opake.dpsgd.clipping import adaptive_clipped_grad
        >>> from opake.dpsgd.noise import gaussian_noise
        >>> from opake.dpsgd.sampling import PoissonSampler
        >>> from opake.distributed import local_shard, sum_gradients, sync
        >>> from opake.random import fold_in, key
        >>>
        >>> # Initialize distributed first (typically via torchrun).
        >>> dist.init_process_group(backend='nccl')
        >>> rank, world_size = dist.get_rank(), dist.get_world_size()
        >>>
        >>> # Create adaptive clipping with the global public expected batch size.
        >>> expected_batch_size = 0.01 * len(dataset)
        >>> grad_fn, clip_state = adaptive_clipped_grad(
        ...     loss_fn,
        ...     key=key(0),
        ...     batch_argnums=(1, 2),
        ...     normalize_by=expected_batch_size,
        ... )
        >>> noise_fn, noise_state = gaussian_noise(noise_multiplier=1.1, key=key(2))
        >>>
        >>> # Poisson sampling: shard per rank, derive a per-rank sampling key.
        >>> shard = local_shard(dataset, rank=rank, world_size=world_size)
        >>> sampler = PoissonSampler(
        ...     shard, sample_rate=0.01, key=fold_in(key(42), rank)
        ... )
        >>> loader = DataLoader(shard, batch_sampler=sampler)
        >>>
        >>> for batch_x, batch_y in loader:
        ...     # Each device: compute clipped gradients and local adaptive state
        ...     grad, clip_state = grad_fn(params, batch_x, batch_y, state=clip_state)
        ...     # Explicit sync step for adaptive clipping state.
        ...     clip_state = sync(clip_state)
        ...
        ...     # Sum clipped gradients across devices
        ...     grad = sum_gradients(grad)
        ...
        ...     # Add noise and update
        ...     noisy_grad, noise_state = noise_fn(grad, noise_state)
        ...     # ... optimizer step

    References:
        Andrew et al., "Differentially Private Learning with Adaptive
        Clipping", NeurIPS 2021. https://arxiv.org/abs/1905.03871
    """
    if return_aux and return_stats:
        raise ConfigurationError(
            *("return_stats cannot be combined with return_aux=True",)
        )

    # Validate parameters
    if isinstance(initial_clipping_norm, PerGroup):
        for gname, val in initial_clipping_norm.values.items():
            if val <= 0:
                raise ConfigurationError(
                    *(
                        f"initial_clipping_norm must be positive for all groups, "
                        f"got {val} for group '{gname}'",
                    )
                )
    elif initial_clipping_norm <= 0:
        raise ConfigurationError(
            *(f"initial_clipping_norm must be positive, got {initial_clipping_norm}",)
        )
    if not 0 < target_quantile < 1:
        raise ConfigurationError(
            *(f"target_quantile must be in (0, 1), got {target_quantile}",)
        )
    if learning_rate <= 0:
        raise ConfigurationError(
            *(f"learning_rate must be positive, got {learning_rate}",)
        )
    if clipping_norm_min <= 0:
        raise ConfigurationError(
            *(f"clipping_norm_min must be positive, got {clipping_norm_min}",)
        )
    if clipping_norm_max <= clipping_norm_min:
        raise ConfigurationError(
            *(
                f"clipping_norm_max ({clipping_norm_max}) must be > clipping_norm_min ({clipping_norm_min})",
            )
        )
    if fraction_noise_std <= 0:
        raise ConfigurationError(
            *(f"fraction_noise_std must be positive, got {fraction_noise_std}",)
        )

    # Normalize argnums/batch_argnums for empty-batch detection.
    argnums_tuple = normalize_to_tuple(argnums)
    batch_argnums_raw = clipped_grad_kwargs.get("batch_argnums", 1)
    batch_argnums_tuple = normalize_to_tuple(batch_argnums_raw)

    # Read static config that the empty-batch path needs to compute max_norms
    # without going through ``clipped_grad`` itself.
    normalize_by = clipped_grad_kwargs.get("normalize_by", 1.0)
    second_moment = bool(clipped_grad_kwargs.get("second_moment", False))

    # Validate the same static args ``clipped_grad`` validates at factory time,
    # so misconfigurations (e.g. ``normalize_by <= 0`` or argnums/batch_argnums
    # collisions) fail at construction here rather than first surfacing on the
    # next non-empty step — keeping empty- and non-empty-batch failure modes
    # consistent.
    _validate_static_args(argnums, batch_argnums_raw, normalize_by)
    if expected_batch_size is None:
        expected_batch_size = clipped_grad_kwargs.get("normalize_by")
    if expected_batch_size is None:
        raise ConfigurationError(
            *("expected_batch_size is required when normalize_by is omitted",)
        )
    if expected_batch_size <= 0:
        raise ConfigurationError(
            *(f"expected_batch_size must be positive, got {expected_batch_size}",)
        )
    output_bound = lambda clipping_norm: clipping_norm / normalize_by  # noqa: E731

    def _output_squared_bound(clipping_norm):
        return (clipping_norm * clipping_norm) / normalize_by

    is_per_group = isinstance(initial_clipping_norm, PerGroup)

    def _empty_num_clipped():
        if is_per_group:
            return dict.fromkeys(sorted(initial_clipping_norm.values), 0.0)
        return 0.0

    def _empty_batch_state(state: AdaptiveClipState) -> AdaptiveClipState:
        """Advance the threshold even when the sampler draws no records."""
        counts = _empty_num_clipped()
        return replace(
            state,
            _current_clipping_norm=state._next_clipping_norm,
            _next_clipping_norm=_updated_clipping_norm(
                state._next_clipping_norm, counts, 0.0, state=state, step=state._step
            ),
            _step=state._step + 1,
            _num_clipped=counts,
            _batch_size=0.0,
        )

    inner_return_aux = return_aux
    inner_fn, _ = clipped_grad(
        loss_fn,
        argnums=argnums,
        has_aux=has_aux,
        clipping_norm=initial_clipping_norm,
        return_aux=inner_return_aux,
        return_stats=not inner_return_aux,
        pre_clipping_transform=pre_clipping_transform,
        **clipped_grad_kwargs,
    )

    def grad_fn(*args, state: AdaptiveClipState, **kwargs):
        """Compute clipped gradients with adaptive threshold.

        Args:
            *args: Positional arguments to pass to the loss function.
            state: Current AdaptiveClipState.
            **kwargs: Keyword arguments to pass to the loss function.

        Returns:
            If return_aux:
                ((grad, aux), new_state)
            Else:
                (grad, new_state)
        """
        with record_function("opake::adaptive_clipped_grad"):
            return _grad_fn_impl(*args, state=state, **kwargs)

    def _grad_fn_impl(*args, state: AdaptiveClipState, **kwargs):
        # Empty batch: zero grads and a noisy threshold update.
        # When ``second_moment=True`` was forwarded to the inner ``clipped_grad``,
        # mirror its empty-batch shape (paired ``SecondMomentClippingOutput``) so
        # downstream noise + optimizer dispatch stays paired across empty and
        # non-empty steps. The squared-stream sensitivity is C^2/normalize_by, the
        # same bound ``clipped_grad`` would attach on a non-empty step.
        if batch_size_from_args(args, batch_argnums_tuple) == 0:
            new_state = _empty_batch_state(state)
            # Build the zero aggregate in the transformed gradient schema.
            zeros = empty_clipped_grads_like(
                args,
                argnums_tuple,
                pre_clipping_transform,
                clipped_grad_kwargs.get("dtype"),
            )
            grads = _empty_batch_grads(
                zeros,
                output_bound(state._next_clipping_norm),
                _output_squared_bound(state._next_clipping_norm),
                second_moment,
            )
            if return_aux:
                empty = torch.empty(0)
                # ``clipping_rate`` presence is schema-driven: None for
                # per-group, like non-empty steps.
                adaptive_aux = AdaptiveClippedGradAux(
                    loss_values=empty,
                    grad_norms=empty,
                    clipped_grad_norms=empty,
                    loss_aux=None,
                    clipping_rate=None if is_per_group else 0.0,
                )
                return (grads, adaptive_aux), new_state
            if return_stats:
                if is_per_group:
                    empty_counts = dict.fromkeys(
                        sorted(initial_clipping_norm.values), 0.0
                    )
                    empty_rates = dict.fromkeys(
                        sorted(initial_clipping_norm.values), 0.0
                    )
                    stats = ClippingStats(
                        num_clipped=empty_counts,
                        clipping_rate=empty_rates,
                        batch_size=0,
                    )
                else:
                    stats = ClippingStats(
                        num_clipped=0.0,
                        clipping_rate=0.0,
                        batch_size=0,
                    )
                return (grads, stats), new_state
            return grads, new_state

        # Compute gradients with the current threshold. The engine transform is
        # constructed once; its tensor-only chunk kernel receives the changing
        # threshold as data so compilation does not specialize on its value.
        user_wants_return_aux = return_aux
        result, _ = inner_fn(
            *args,
            state=_RuntimeClipState(state._next_clipping_norm),
            **kwargs,
        )

        # Extract gradients and auxiliary output
        aux = None
        stats: ClippingStats | None = None
        if isinstance(result, tuple):
            if inner_return_aux:
                grads, aux = result
                grad_norms = aux.grad_norms
            else:
                grads, stats = result
                grad_norms = None
        else:
            grads = result
            grad_norms = None

        batch_size = (
            aux.batch_size
            if aux is not None
            else (stats.batch_size if stats is not None else 0)
        )

        if is_per_group and (grad_norms is not None or stats is not None):
            current_pg = state._next_clipping_norm
            if stats is not None and isinstance(stats.num_clipped, dict):
                num_clipped = stats.num_clipped
            elif aux is not None and aux.group_norms is not None:
                num_clipped = {
                    name: float(
                        (aux.group_norms[name] > current_pg.values[name]).sum().item()
                    )
                    for name in sorted(current_pg.values)
                }
            else:
                num_clipped = _empty_num_clipped()

            new_clipping_norm = _updated_clipping_norm(
                current_pg,
                num_clipped,
                float(batch_size),
                state=state,
                step=state._step,
            )
        elif grad_norms is not None or stats is not None:
            # --- Scalar adaptive path ---
            if stats is not None and isinstance(stats.num_clipped, float):
                num_clipped = stats.num_clipped
            else:
                num_clipped = float(
                    (grad_norms > state._next_clipping_norm).sum().item()
                )

            new_clipping_norm = _updated_clipping_norm(
                state._next_clipping_norm,
                num_clipped,
                float(batch_size),
                state=state,
                step=state._step,
            )
        else:
            new_clipping_norm = state._next_clipping_norm
            num_clipped = _empty_num_clipped()

        new_state = replace(
            state,
            _current_clipping_norm=state._next_clipping_norm,
            _next_clipping_norm=new_clipping_norm,
            _step=state._step + 1,
            _num_clipped=num_clipped,
            _batch_size=float(batch_size),
        )

        if user_wants_return_aux:
            adaptive_aux = AdaptiveClippedGradAux(
                loss_values=aux.loss_values if aux is not None else None,
                grad_norms=aux.grad_norms if aux is not None else None,
                clipped_grad_norms=(
                    aux.clipped_grad_norms if aux is not None else None
                ),
                loss_aux=aux.loss_aux if aux is not None else None,
                clipping_rate=aux.clipping_rate if aux is not None else None,
                batch_size=batch_size,
                group_norms=aux.group_norms if aux is not None else None,
            )
            return (grads, adaptive_aux), new_state

        if return_stats:
            assert stats is not None
            return (grads, stats), new_state
        return grads, new_state

    # Create initial state
    initial_state = AdaptiveClipState(
        _current_clipping_norm=initial_clipping_norm,
        _next_clipping_norm=initial_clipping_norm,
        _step=0,
        _rng_key=key,
        _fraction_noise_std=fraction_noise_std,
        _expected_batch_size=expected_batch_size,
        _learning_rate=learning_rate,
        _target_quantile=target_quantile,
        _clipping_norm_min=clipping_norm_min,
        _clipping_norm_max=clipping_norm_max,
        _num_clipped=_empty_num_clipped(),
        _batch_size=0,
    )

    return grad_fn, initial_state

`auto_clipped_grad(loss_fn, argnums=0, has_aux=False, *, R=1.0, gamma=_DEFAULT_GAMMA, normalize_by=1.0, batch_argnums=1, return_aux=False, return_stats=False, pre_clipping_transform=lambda x: x, microbatch_size=None, dtype=None, second_moment=False, _chunk_compiler=None)` ¶

Create a function that computes the sum of AUTO-S scaled per-example gradients.

AUTO-S (Bu et al., 2023) replaces the standard clipping threshold with a per-example scaling

.. math::

\tilde g_i = R \cdot g_i / (\lVert g_i \rVert + \gamma)

so every per-example gradient has L2 norm at most R by construction. There is no threshold to tune; the effective step size is absorbed into the optimizer learning rate.

Parameters:

Name	Type	Description	Default
`loss_fn`	`Callable`	Scalar loss function. If `has_aux`, returns `(scalar, loss_aux)`.	required
`argnums`	`int \| tuple[int, ...]`	Which argument(s) to differentiate w.r.t.	`0`
`has_aux`	`bool`	If True, `loss_fn` returns `(scalar, loss_aux)`.	`False`
`R`	`float \| PerGroup`	Sensitivity max_norm (default 1.0). When `PerGroup`, each parameter group is scaled independently to its own `R_k`.	`1.0`
`gamma`	`float`	Denominator stabilizer :math:`\gamma` (default 0.01, strictly positive).	`_DEFAULT_GAMMA`
`normalize_by`	`float`	Divisor applied to the summed gradients (set to expected batch size for averaged gradients).	`1.0`
`batch_argnums`	`int \| tuple[int, ...]`	Which arguments have a batch dimension.	`1`
`return_aux`	`bool`	If True, returns per-example diagnostics as :class:`AutoClippedGradAux`.	`False`
`pre_clipping_transform`	`Callable`	An optional function to apply to the per-example gradients before AUTO-S scaling. The function should consume the gradient pytree for a single example and return a new pytree (possibly with different structure). Can be used to e.g., scale the leaves of the pytree to accommodate preconditioner clipping. Does not affect the sensitivity guarantee. Default is identity function.	`lambda x: x`
`microbatch_size`	`int \| None`	Process the batch in chunks of this size.	`None`
`dtype`	`Any`	Optional accumulation dtype for the summed gradient.	`None`
`second_moment`	`bool`	If True, also accumulate the per-example sum of element-wise squared scaled gradients and return a :class:`~opake.types.SecondMomentClippingOutput` carrying both streams. Privacy accounting is unchanged (`gaussian(noise_multiplier)` for DP-SGD or any standard MF mechanism for DP-FTRL); the sensitivity-proportional joint Mahalanobis allocation gives the paired release the same PLD as a single first-moment release.	`False`
`return_stats`	`bool`	If True, return :class:`ClippingStats` beside the gradient. Cannot be combined with `return_aux`.	`False`

Returns:

Type	Description
`Callable`	`(grad_fn, state)`. `grad_fn` has the signature
`AutoClipState`	`(args, state, *kwargs) -> (grad, state)`, or
`tuple[Callable, AutoClipState]`	`(args, state, *kwargs) -> ((grad, aux), state)` when
`tuple[Callable, AutoClipState]`	`return_aux=True`. When `return_stats=True`, it returns
`tuple[Callable, AutoClipState]`	`((grad, stats), state)`.

Formal guarantee

Under add/remove or zero-out DP, the L2 sensitivity of the summed gradients is the returned ClippedPytree.max_norm metadata — a constant independent of the input data. Per-step Gaussian noise (gaussian_noise, …) reads that value each step; PLD composition does not require max_norm to be identical across steps when the accountant models step-varying sensitivity (for example adaptive clipping with adaclip). mf_gaussian_noise's matrix-factorization correlated noise does require max_norm to stay fixed for the whole run — the dispatcher latches the first-call bound — and AUTO-S satisfies that because R and normalize_by do not drift. AUTO-S scaling is per-example and adds no privacy cost beyond gaussian(noise_multiplier) for the gradient release.

Example

import torch from opake.api.engine.clipping import auto_clipped_grad def loss_fn(params, x, y): ... return ((x @ params - y) ** 2).mean() grad_fn, state = auto_clipped_grad( ... loss_fn, argnums=0, batch_argnums=(1, 2), ... R=1.0, normalize_by=32, ... ) params = torch.randn(10) batch_x = torch.randn(32, 10) batch_y = torch.randn(32) grads, state = grad_fn(params, batch_x, batch_y, state=state)

The returned ClippedPytree carries max_norm = R / normalize_by (constant across steps), so it composes with mf_gaussian_noise (DP-FTRL) and gaussian_noise (DP-SGD) the same way clipped_grad does at the same bound.

References

Bu, Wang, Zha, Karypis. "Automatic Clipping: Differentially Private Deep Learning Made Easier and Stronger." NeurIPS 2023. https://arxiv.org/abs/2206.07136

Source code in packages/opake-engine/src/opake/api/engine/clipping/_auto.py

def auto_clipped_grad(
    loss_fn: Callable,
    argnums: int | tuple[int, ...] = 0,
    has_aux: bool = False,
    *,
    R: float | PerGroup = 1.0,
    gamma: float = _DEFAULT_GAMMA,
    normalize_by: float = 1.0,
    batch_argnums: int | tuple[int, ...] = 1,
    return_aux: bool = False,
    return_stats: bool = False,
    pre_clipping_transform: Callable = lambda x: x,
    microbatch_size: int | None = None,
    dtype: Any = None,
    second_moment: bool = False,
    _chunk_compiler: Callable | None = None,
) -> tuple[Callable, AutoClipState]:
    r"""Create a function that computes the sum of AUTO-S scaled per-example gradients.

    AUTO-S (Bu et al., 2023) replaces the standard clipping threshold with
    a per-example scaling

    .. math::

        \tilde g_i = R \cdot g_i / (\lVert g_i \rVert + \gamma)

    so every per-example gradient has L2 norm at most ``R`` by
    construction.  There is no threshold to tune; the effective step size
    is absorbed into the optimizer learning rate.

    Args:
        loss_fn: Scalar loss function.  If ``has_aux``, returns
            ``(scalar, loss_aux)``.
        argnums: Which argument(s) to differentiate w.r.t.
        has_aux: If True, ``loss_fn`` returns ``(scalar, loss_aux)``.
        R: Sensitivity max_norm (default 1.0).  When ``PerGroup``, each
            parameter group is scaled independently to its own ``R_k``.
        gamma: Denominator stabilizer :math:`\gamma` (default 0.01,
            strictly positive).
        normalize_by: Divisor applied to the summed gradients (set to
            expected batch size for averaged gradients).
        batch_argnums: Which arguments have a batch dimension.
        return_aux: If True, returns per-example diagnostics as
            :class:`AutoClippedGradAux`.
        pre_clipping_transform: An optional function to apply to the
            per-example gradients before AUTO-S scaling. The function
            should consume the gradient pytree for a single example
            and return a new pytree (possibly with different structure).
            Can be used to e.g., scale the leaves of the pytree to
            accommodate preconditioner clipping. Does not affect the
            sensitivity guarantee. Default is identity function.
        microbatch_size: Process the batch in chunks of this size.
        dtype: Optional accumulation dtype for the summed gradient.
        second_moment: If True, also accumulate the per-example sum of
            element-wise squared scaled gradients and return a
            :class:`~opake.types.SecondMomentClippingOutput` carrying
            both streams.  Privacy accounting is unchanged
            (``gaussian(noise_multiplier)`` for DP-SGD or any standard MF
            mechanism for DP-FTRL); the sensitivity-proportional joint
            Mahalanobis allocation gives the paired release the same PLD
            as a single first-moment release.
        return_stats: If True, return :class:`ClippingStats` beside the
            gradient. Cannot be combined with ``return_aux``.

    Returns:
        ``(grad_fn, state)``.  ``grad_fn`` has the signature
        ``(*args, state, **kwargs) -> (grad, state)``, or
        ``(*args, state, **kwargs) -> ((grad, aux), state)`` when
        ``return_aux=True``. When ``return_stats=True``, it returns
        ``((grad, stats), state)``.

    Formal guarantee:
        Under add/remove or zero-out DP, the L2 sensitivity of the
        summed gradients is the returned ``ClippedPytree.max_norm`` metadata —
        a constant independent of the input data.  Per-step Gaussian noise
        (``gaussian_noise``, …) reads that value
        each step; PLD composition does not require ``max_norm`` to be
        identical across steps when the accountant models step-varying
        sensitivity (for example adaptive clipping with ``adaclip``).
        ``mf_gaussian_noise``'s matrix-factorization correlated noise *does* require
        ``max_norm`` to stay fixed for the whole run — the dispatcher latches
        the first-call bound — and AUTO-S satisfies that because ``R`` and
        ``normalize_by`` do not drift.  AUTO-S scaling is per-example and adds
        no privacy cost beyond ``gaussian(noise_multiplier)`` for the
        gradient release.

    Example:
        >>> import torch
        >>> from opake.api.engine.clipping import auto_clipped_grad
        >>> def loss_fn(params, x, y):
        ...     return ((x @ params - y) ** 2).mean()
        >>> grad_fn, state = auto_clipped_grad(
        ...     loss_fn, argnums=0, batch_argnums=(1, 2),
        ...     R=1.0, normalize_by=32,
        ... )
        >>> params = torch.randn(10)
        >>> batch_x = torch.randn(32, 10)
        >>> batch_y = torch.randn(32)
        >>> grads, state = grad_fn(params, batch_x, batch_y, state=state)

    The returned ``ClippedPytree`` carries ``max_norm = R / normalize_by``
    (constant across steps), so it composes with ``mf_gaussian_noise`` (DP-FTRL)
    and ``gaussian_noise`` (DP-SGD) the same way ``clipped_grad`` does at
    the same bound.

    References:
        Bu, Wang, Zha, Karypis.  "Automatic Clipping: Differentially
        Private Deep Learning Made Easier and Stronger." NeurIPS 2023.
        https://arxiv.org/abs/2206.07136
    """
    _validate_auto_params(R, gamma)

    scale_fn = _make_auto_scale_fn(R, gamma)
    builtin_scale = _BuiltinScale(kind="auto_s", gamma=gamma)
    inner_fn, _ = clipped_grad(
        loss_fn,
        argnums=argnums,
        has_aux=has_aux,
        clipping_norm=R,
        normalize_by=normalize_by,
        batch_argnums=batch_argnums,
        return_aux=return_aux,
        return_stats=return_stats,
        pre_clipping_transform=pre_clipping_transform,
        microbatch_size=microbatch_size,
        dtype=dtype,
        second_moment=second_moment,
        _scale_fn=scale_fn,
        _builtin_scale=builtin_scale,
        _chunk_compiler=_chunk_compiler,
    )

    state = AutoClipState()

    if not return_aux:

        def grad_fn(*args, state, **kwargs):
            result, _ = inner_fn(*args, state=None, **kwargs)
            return result, state

        return grad_fn, state

    def grad_fn(*args, state, **kwargs):
        inner_result, _ = inner_fn(*args, state=None, **kwargs)
        grads, grad_aux = inner_result
        auto_aux = AutoClippedGradAux(
            loss_values=grad_aux.loss_values,
            grad_norms=grad_aux.grad_norms,
            clipped_grad_norms=grad_aux.clipped_grad_norms,
            loss_aux=grad_aux.loss_aux,
            clipping_rate=grad_aux.clipping_rate,
            batch_size=grad_aux.batch_size,
            group_norms=grad_aux.group_norms,
        )
        return (grads, auto_aux), state

    return grad_fn, state

`clipped_grad(loss_fn, argnums=0, has_aux=False, *, clipping_norm, normalize_by=1.0, batch_argnums=1, return_aux=False, return_stats=False, second_moment=False, pre_clipping_transform=lambda x: x, microbatch_size=None, dtype=None, compute_dtype=None, _scale_fn=None, _builtin_scale=None, _chunk_compiler=None)` ¶

Create a function to compute the sum of clipped gradients of loss_fn.

This function acts as a transformation similar to torch.func.grad, but with added functionality for gradient clipping applied on a per-example (or per-group) basis before summation. It computes the gradient of loss_fn with respect to argnums, calculates the L2 norm of the gradient for each example slice along the first axis of the batch_argnums args, clips each per-example gradient to have a norm of at most clipping_norm, and finally sums these clipped gradients.

Non-grad outputs of the returned function (aux) may optionally be returned by setting return_aux=True. These outputs are per-example, and hence have a batch axis. It is up to the caller to handle these as necessary.

Example Usage

import torch from opake.dpsgd.clipping import clipped_grad f = lambda param, data: 0.5 * ((data - param) ** 2).mean() g, clip_state = clipped_grad(f, clipping_norm=float('inf')) result, clip_state = g(torch.tensor(3.0), torch.tensor([0.0, 7.0, -2.0]), state=clip_state) result tensor(1.3333)

Example Usage (with Auxiliary Output): >>> g, clip_state = clipped_grad( ... f, clipping_norm=float('inf'), return_aux=True ... ) >>> (_, aux), clip_state = g(torch.tensor(3.0), torch.tensor([0.0, 7.0, -2.0]), state=clip_state) >>> aux.loss_values tensor([4.5000, 8.0000, 12.5000]) >>> aux.grad_norms tensor([3., 4., 5.])

Formal Guarantees

For the gradient output: The L2 sensitivity of the returned function with respect to the batch arguments (specified by batch_argnums) under add/remove or zero-out differential privacy definitions is guaranteed to be clipping_norm. Under replace-one DP, the sensitivity is doubled (2 * clipping_norm). All auxiliary outputs (loss_values, grad_norms) are per-example. This function guarantees that per-example outputs only depend on the data for the same example. This allows maximum flexibility for the caller to aggregate these as desired (possibly with a DP mean, median, quantile, or histogram mechanism).

Parameters:

Name	Type	Description	Default
`loss_fn`	`Callable`	The loss function to be differentiated, which should return a scalar value. If `has_aux` is True, it should return a tuple `(value, loss_aux)`.	required
`argnums`	`int \| tuple[int, ...]`	Specifies which argument(s) of `loss_fn` to differentiate with respect to. Can be an integer or a sequence of integers. These arguments should not have a batch dimension.	`0`
`has_aux`	`bool`	If True, `loss_fn` is expected to return a tuple `(value, loss_aux)`. The auxiliary data `loss_aux` will be returned by the transformed function. Exercise caution when using this as no DP sensitivity guarantees are provided for the auxiliary data.	`False`
`clipping_norm`	`float \| PerGroup`	The maximum L2 norm for each per-example gradient. Gradients with a norm larger than this value will be scaled down.	required
`normalize_by`	`float`	Divide the clipped sum by this constant before returning. Set to expected batch size to produce averaged gradients with sensitivity = clipping_norm / normalize_by.	`1.0`
`batch_argnums`	`int \| tuple[int, ...]`	Specifies which argument(s) of `loss_fn` contain the batch dimension (usually the data and labels). Can be an integer or a sequence of integers. All arguments specified here must have the same size along their first dimension (the batch dimension). The default value of 1 assumes the signature of loss_fn is `loss_fn(params, batch)`.	`1`
`return_aux`	`bool`	If True, the transformed function will also return a per-example aux dataclass containing loss values, gradient norms, and loss aux.	`False`
`return_stats`	`bool`	If True, return aggregate clipping statistics including the number and fraction of clipped examples and the batch size. Cannot be combined with `return_aux`.	`False`
`pre_clipping_transform`	`Callable`	An optional function to apply to the per-example gradients before clipping. The function should consume the gradient pytree for a single example and return a new pytree (possibly with different structure). Can be used to e.g., scale the leaves of the pytree to accommodate preconditioner clipping. Account separately for any data-dependent control flow or telemetry. Default is identity function.	`lambda x: x`
`microbatch_size`	`int \| None`	If set, the batch is split up into microbatches of this size for memory-efficient processing. Processes each microbatch separately and accumulates results without materializing the full batch of gradients. Set this to reduce peak memory usage at the cost of slightly slower computation.	`None`
`dtype`	`dtype \| None`	Optional dtype for the returned gradient. If None, the dtype will be the same as the dtypes of the gradient function. Can be useful to avoid overflow issues when using low-precision dtypes as the returned function computes a sum over a potentially large batch.	`None`
`compute_dtype`	`dtype \| None`	Internal reduction and accumulation dtype. `None` (default) auto-promotes bf16/fp16 to float32 for numerical stability; an explicit dtype selects the reduction precision. Leaf scaling and microbatch accumulation use the wider of this dtype and the leaf or output storage dtype, so explicit lower precision never narrows a wider value. Independent of `dtype` (which controls the output dtype).	`None`
`second_moment`	`bool`	Whether to accumulate the clipped-gradient second moment required by DP-FTRL noise mechanisms.	`False`

Returns: Tuple of (:class:ClippedGradFn, clip_state) where: - clipped_grad_fn: A function that computes the sum of clipped per-example gradients. Call signature: clipped_grads, new_state = clipped_grad_fn(..., state=clip_state) If auxiliary outputs are requested, returns: (clipped_grads, grad_aux), new_state If return_stats=True, returns (clipped_grads, stats), new_state. - clip_state: Initial FixedClipState containing sensitivity information

The grad_aux output (when requested) is a ClippedGradAux dataclass with fields:
    - loss_values: Per-example function values (if return_aux=True)
    - grad_norms: Per-example gradient norms (if return_aux=True)
    - loss_aux: Per-example auxiliary data (if has_aux=True)

Source code in packages/opake-engine/src/opake/api/engine/clipping/_clipped_grad.py

def clipped_grad(
    loss_fn: Callable,
    argnums: int | tuple[int, ...] = 0,
    has_aux: bool = False,
    *,
    clipping_norm: float | PerGroup,
    normalize_by: float = 1.0,
    batch_argnums: int | tuple[int, ...] = 1,
    return_aux: bool = False,
    return_stats: bool = False,
    second_moment: bool = False,
    pre_clipping_transform: Callable = lambda x: x,
    microbatch_size: int | None = None,
    dtype: torch.dtype | None = None,
    compute_dtype: torch.dtype | None = None,
    _scale_fn: Callable | None = None,
    _builtin_scale: _BuiltinScale | None = None,
    _chunk_compiler: Callable | None = None,
) -> tuple[ClippedGradFn, FixedClipState]:
    """Create a function to compute the sum of clipped gradients of loss_fn.

    This function acts as a transformation similar to `torch.func.grad`, but with added
    functionality for gradient clipping applied on a per-example (or per-group)
    basis before summation. It computes the gradient of `loss_fn` with respect to
    `argnums`, calculates the L2 norm of the gradient for each example slice
    along the first axis of the `batch_argnums` args, clips each per-example
    gradient to have a norm of at most `clipping_norm`, and finally sums these
    clipped gradients.

    Non-grad outputs of the returned function (aux) may optionally be returned
    by setting `return_aux=True`. These outputs are per-example, and hence have
    a batch axis. It is up to the caller to handle these as necessary.

    Example Usage:
        >>> import torch
        >>> from opake.dpsgd.clipping import clipped_grad
        >>> f = lambda param, data: 0.5 * ((data - param) ** 2).mean()
        >>> g, clip_state = clipped_grad(f, clipping_norm=float('inf'))
        >>> result, clip_state = g(torch.tensor(3.0), torch.tensor([0.0, 7.0, -2.0]), state=clip_state)
        >>> result
        tensor(1.3333)

    Example Usage (with Auxiliary Output):
        >>> g, clip_state = clipped_grad(
        ...     f, clipping_norm=float('inf'), return_aux=True
        ... )
        >>> (_, aux), clip_state = g(torch.tensor(3.0), torch.tensor([0.0, 7.0, -2.0]), state=clip_state)
        >>> aux.loss_values
        tensor([4.5000, 8.0000, 12.5000])
        >>> aux.grad_norms
        tensor([3., 4., 5.])

    Formal Guarantees:
        For the gradient output:
          The L2 sensitivity of the returned function with respect to the batch
          arguments (specified by `batch_argnums`) under add/remove or zero-out
          differential privacy definitions is guaranteed to be `clipping_norm`.
          Under replace-one DP, the sensitivity is doubled (2 * `clipping_norm`).
        All auxiliary outputs (loss_values, grad_norms) are per-example. This
          function guarantees that per-example outputs only depend on the data for the
          same example. This allows maximum flexibility for the caller to aggregate
          these as desired (possibly with a DP mean, median, quantile, or histogram
          mechanism).

    Args:
        loss_fn: The loss function to be differentiated, which should return a scalar
            value. If `has_aux` is True, it should return a tuple `(value, loss_aux)`.
        argnums: Specifies which argument(s) of `loss_fn` to differentiate with respect
            to. Can be an integer or a sequence of integers. These arguments should
            *not* have a batch dimension.
        has_aux: If True, `loss_fn` is expected to return a tuple `(value, loss_aux)`.
            The auxiliary data `loss_aux` will be returned by the transformed function.
            Exercise caution when using this as no DP sensitivity guarantees are
            provided for the auxiliary data.
        clipping_norm: The maximum L2 norm for each per-example gradient. Gradients
            with a norm larger than this value will be scaled down.
        normalize_by: Divide the clipped sum by this constant before returning.
            Set to expected batch size to produce averaged gradients with
            sensitivity = clipping_norm / normalize_by.
        batch_argnums: Specifies which argument(s) of `loss_fn` contain the batch
            dimension (usually the data and labels). Can be an integer or a sequence
            of integers. All arguments specified here must have the same size along
            their first dimension (the batch dimension). The default value of 1 assumes
            the signature of loss_fn is `loss_fn(params, batch)`.
        return_aux: If True, the transformed function will also return a per-example
            aux dataclass containing loss values, gradient norms, and loss aux.
        return_stats: If True, return aggregate clipping statistics including
            the number and fraction of clipped examples and the batch size.
            Cannot be combined with ``return_aux``.
        pre_clipping_transform: An optional function to apply to the per-example
            gradients before clipping. The function should consume the gradient pytree
            for a single example and return a new pytree (possibly with different
            structure). Can be used to e.g., scale the leaves of the pytree to
            accommodate preconditioner clipping. Account separately for any
            data-dependent control flow or telemetry. Default is identity function.
        microbatch_size: If set, the batch is split up into microbatches of this
            size for memory-efficient processing. Processes each microbatch separately
            and accumulates results without materializing the full batch of gradients.
            Set this to reduce peak memory usage at the cost of slightly slower computation.
        dtype: Optional dtype for the returned gradient. If None, the dtype will be
            the same as the dtypes of the gradient function. Can be useful to avoid
            overflow issues when using low-precision dtypes as the returned function
            computes a sum over a potentially large batch.
        compute_dtype: Internal reduction and accumulation dtype. ``None``
            (default) auto-promotes bf16/fp16 to float32 for numerical stability;
            an explicit dtype selects the reduction precision. Leaf scaling and
            microbatch accumulation use the wider of this dtype and the leaf or
            output storage dtype, so explicit lower precision never narrows a wider
            value. Independent of ``dtype`` (which controls the *output* dtype).
        second_moment: Whether to accumulate the clipped-gradient second
            moment required by DP-FTRL noise mechanisms.
    Returns:
        Tuple of (:class:`ClippedGradFn`, clip_state) where:
        - clipped_grad_fn: A function that computes the sum of clipped per-example gradients.
          Call signature: clipped_grads, new_state = clipped_grad_fn(..., state=clip_state)
          If auxiliary outputs are requested, returns: (clipped_grads, grad_aux), new_state
          If ``return_stats=True``, returns ``(clipped_grads, stats), new_state``.
        - clip_state: Initial FixedClipState containing sensitivity information

        The grad_aux output (when requested) is a ClippedGradAux dataclass with fields:
            - loss_values: Per-example function values (if return_aux=True)
            - grad_norms: Per-example gradient norms (if return_aux=True)
            - loss_aux: Per-example auxiliary data (if has_aux=True)
    """
    _validate_static_args(argnums, batch_argnums, normalize_by)
    if return_aux and return_stats:
        raise ConfigurationError(
            *("return_stats cannot be combined with return_aux=True",)
        )

    argnums_tuple = normalize_to_tuple(argnums)
    batch_argnums_tuple = normalize_to_tuple(batch_argnums)
    loss_fn = normalize_fun_to_return_aux(loss_fn, has_aux)

    output_max_norm = clipping_norm / normalize_by
    output_squared_max_norm = (
        (clipping_norm * clipping_norm) / normalize_by if second_moment else None
    )

    def _empty_batch_response(args, state):
        """Return zero gradients and empty diagnostics without vmap."""
        zeros = empty_clipped_grads_like(
            args, argnums_tuple, pre_clipping_transform, dtype
        )
        if second_moment:
            from opake.api.engine.types import SecondMomentClippingOutput

            grads = SecondMomentClippingOutput(
                grads=clipped(zeros, max_norm=output_max_norm),
                squared_grads=clipped(
                    tree_map(
                        lambda x: (
                            torch.zeros_like(x) if isinstance(x, torch.Tensor) else x
                        ),
                        zeros,
                    ),
                    max_norm=output_squared_max_norm,
                ),
            )
        else:
            grads = clipped(zeros, max_norm=output_max_norm)
        if return_aux:
            empty = torch.empty(0)
            # Per-group steps report ``clipping_rate=None`` (rates live in
            # ``group_norms``), empty draws included, so ``sync(aux)`` matches.
            grad_aux = ClippedGradAux(
                loss_values=empty if return_aux else None,
                grad_norms=empty,
                clipped_grad_norms=empty,
                loss_aux=None,
                clipping_rate=None if isinstance(clipping_norm, PerGroup) else 0.0,
                batch_size=0,
            )
            return (grads, grad_aux), state
        if return_stats:
            from opake.api.engine.clipping._clipped_fun import ClippingStats

            if isinstance(clipping_norm, PerGroup):
                group_names = sorted(clipping_norm.values)
                empty_counts = dict.fromkeys(group_names, 0.0)
                empty_rates = dict.fromkeys(group_names, 0.0)
                stats = ClippingStats(
                    num_clipped=empty_counts,
                    clipping_rate=empty_rates,
                    batch_size=0,
                )
            else:
                stats = ClippingStats(num_clipped=0.0, clipping_rate=0.0, batch_size=0)
            return (grads, stats), state
        return grads, state

    # Use PyTorch's grad_and_value (returns (grad, value) or (grad, (value, aux)))
    grad_and_value_fn = grad_and_value(loss_fn, argnums=argnums, has_aux=True)

    def grad_fn(*args, **kwargs):
        # Keep the graph only when an enclosing transform differentiates it.
        values_only = not under_differentiating_transform(when_compiling=True)
        with torch.no_grad() if values_only else contextlib.nullcontext():
            grad, value_and_aux = grad_and_value_fn(*args, **kwargs)
        result = pre_clipping_transform(grad)
        if return_aux:
            # Return dict aux from per-example grad_fn; clipping-related norms are
            # produced by clipped_fun to avoid duplicating norm computation logic.
            # PyTorch vmap cannot handle namedtuples with None values when out_dims != None
            aux_dict = {}
            if return_aux:
                aux_dict["values"] = value_and_aux[0]
            if return_aux and has_aux:
                aux_dict["value_aux"] = value_and_aux[1]
            return result, aux_dict
        return result

    clipped_grad_fn, clip_state = clipped_fun(
        grad_fn,
        has_aux=return_aux,
        batch_argnums=batch_argnums,
        clipping_norm=clipping_norm,
        normalize_by=normalize_by,
        return_aux=return_aux,
        return_stats=return_stats,
        second_moment=second_moment,
        microbatch_size=microbatch_size,
        dtype=dtype,
        compute_dtype=compute_dtype,
        _scale_fn=_scale_fn,
        _builtin_scale=_builtin_scale,
        _chunk_compiler=_chunk_compiler,
    )

    if return_stats:

        def grad_fn_wrapper(*args, state, **kwargs):
            if batch_size_from_args(args, batch_argnums_tuple) == 0:
                return _empty_batch_response(args, state)
            with record_function("opake::clipped_grad"):
                (clipped_grads, stats), returned_state = clipped_grad_fn(
                    *args, state=state, **kwargs
                )
            return (clipped_grads, stats), returned_state

        return grad_fn_wrapper, clip_state

    # clipped_grad_fn is now a callable, clip_state is a FixedClipState
    # Wrap the result to convert ClippedFunAux to ClippedGradAux
    if not return_aux:
        # No aux, return wrapped directly with state-passing signature
        def grad_fn_wrapper(*args, state, **kwargs):
            if batch_size_from_args(args, batch_argnums_tuple) == 0:
                return _empty_batch_response(args, state)

            with record_function("opake::clipped_grad"):
                (result, returned_state) = clipped_grad_fn(*args, state=state, **kwargs)
            return result, returned_state

        return grad_fn_wrapper, clip_state
    else:
        # Need to convert ClippedFunAux to ClippedGradAux
        def grad_fn_wrapper(*args, state, **kwargs):
            if batch_size_from_args(args, batch_argnums_tuple) == 0:
                return _empty_batch_response(args, state)

            with record_function("opake::clipped_grad"):
                (clipped_grads, aux), returned_state = clipped_grad_fn(
                    *args, state=state, **kwargs
                )
            grad_aux = ClippedGradAux(
                loss_values=aux.values,
                grad_norms=aux.norms,
                clipped_grad_norms=aux.clipped_norms,
                loss_aux=aux.value_aux,
                clipping_rate=aux.clipping_rate,
                batch_size=aux.batch_size,
                group_norms=aux.group_norms,
            )
            return (clipped_grads, grad_aux), returned_state

        return grad_fn_wrapper, clip_state

`per_group(params, /, patterns=None, *, fallback=None, allow_unused_patterns=False, **kwargs)` ¶

Construct PerGroup from parameter leaf paths and substring patterns.

Each pattern is a substring matched against the dotted display form of each leaf :data:~opake.pytree.ParamPath (see :func:~opake.pytree.param_path_display). Leaves whose display path contains the substring are assigned to that group. Every leaf must match exactly one pattern (error on 0 or 2+).

params may be any tensor pytree: flat named_parameters dicts (paths are one-segment (name,)), nested dicts, lists, or tuples. The returned :class:~opake.types.PerGroup is compiled against those optree paths so clip / noise / optimizers look up the same identity.

The patterns dict arg merges with **kwargs to support keys containing dots (which cannot be used as keyword arguments)::

per_group(params, **{"layers.0": 0.5, "layers.1": 1.0})
# equivalent to:
per_group(params, patterns={"layers.0": 0.5, "layers.1": 1.0})

The fallback parameter assigns a value to any parameter that does not match an explicit pattern. Without fallback, unmatched parameters raise ValueError.

By default, every explicit pattern must match at least one parameter. Set allow_unused_patterns=True when intentionally reusing a configuration across models with different parameter names.

Parameters:

Name	Type	Description	Default
`params`	`Any`	Parameter pytree (flat or nested). Leaf paths are matched against patterns via their dotted display form.	required
`patterns`	`dict[str, float] \| None`	Optional dict of `{pattern: value}` pairs. Merged with `**kwargs`.	`None`
`fallback`	`float \| None`	Optional value for parameters that don't match any explicit pattern. Unmatched params are assigned to a group named `"fallback"`.	`None`
`allow_unused_patterns`	`bool`	Whether to allow explicit patterns that do not match any parameter. Defaults to `False`.	`False`
`**kwargs`	`float`	Pattern-value pairs where the kwarg name is the substring pattern and the value is the per-group value (e.g. clipping norm).	`{}`

Returns:

Type	Description
`PerGroup`	class:`~opake.types.PerGroup` with pre-resolved
`PerGroup`	path-to-group assignments.

`fun` ¶

DP-SGD clipping power-user helpers.

:func:clipped_fun / :func:auto_clipped_fun — clip + sum any per-example function output (PyTree).
:func:clip_pytree / :func:auto_scale_pytree — clip / AUTO-S an already-batched pytree.

`auto_clipped_fun(fun, has_aux=False, *, batch_argnums=0, R=1.0, gamma=_DEFAULT_GAMMA, normalize_by=1.0, return_aux=False, return_stats=False, microbatch_size=None, dtype=None, _chunk_compiler=None)` ¶

Transform a function so each per-example output is scaled to clipped norm via AUTO-S.

Mirrors :func:clipped_fun but replaces threshold-based clipping with AUTO-S automatic scaling R \cdot v / (\|v\| + \gamma). The output of the returned function is the sum of the scaled per-example outputs.

Parameters:

Name	Type	Description	Default
`fun`	`Callable`	The per-example function to be transformed.	required
`has_aux`	`bool`	If True, `fun` returns `(value, aux)`; only `value` is scaled and aggregated.	`False`
`batch_argnums`	`int \| tuple[int, ...]`	Which arguments have a batch dimension.	`0`
`R`	`float \| PerGroup`	Sensitivity max_norm for the scaled output. When `PerGroup`, each group is scaled independently.	`1.0`
`gamma`	`float`	Denominator stabilizer :math:`\gamma` (default 0.01).	`_DEFAULT_GAMMA`
`normalize_by`	`float`	Divisor applied to the scaled sum.	`1.0`
`return_aux`	`bool`	If True, the returned callable returns an :class:`AutoClippedFunAux` alongside the summed value.	`False`
`return_stats`	`bool`	If True, return :class:`ClippingStats` alongside the summed value. Cannot be combined with `return_aux`.	`False`
`microbatch_size`	`int \| None`	Process the batch in chunks of this size.	`None`
`dtype`	`Any`	Optional accumulation dtype for the sum.	`None`

Returns:

Type	Description
`Callable`	`(auto_fn, state)`. `auto_fn` has the signature
`AutoClipState`	`(args, state, *kwargs) -> (value, state)` or
`tuple[Callable, AutoClipState]`	`(args, state, *kwargs) -> ((value, aux), state)` when
`tuple[Callable, AutoClipState]`	`return_aux=True`.

Raises:

Type	Description
`ConfigurationError`	If the batch is empty.

Formal guarantee

Under add/remove or zero-out DP, the L2 sensitivity of the first output with respect to the batch arguments is the returned ClippedPytree.max_norm metadata. The bound is constant and data-independent, so per-step Gaussian calibration is correct. mf_gaussian_noise additionally requires that max_norm not drift across steps; AUTO-S satisfies that latch because R and normalize_by are fixed.

Source code in packages/opake-engine/src/opake/api/engine/clipping/_auto.py

def auto_clipped_fun(
    fun: Callable,
    has_aux: bool = False,
    *,
    batch_argnums: int | tuple[int, ...] = 0,
    R: float | PerGroup = 1.0,
    gamma: float = _DEFAULT_GAMMA,
    normalize_by: float = 1.0,
    return_aux: bool = False,
    return_stats: bool = False,
    microbatch_size: int | None = None,
    dtype: Any = None,
    _chunk_compiler: Callable | None = None,
) -> tuple[Callable, AutoClipState]:
    r"""Transform a function so each per-example output is scaled to clipped norm via AUTO-S.

    Mirrors :func:`clipped_fun` but replaces threshold-based clipping with
    AUTO-S automatic scaling ``R \cdot v / (\|v\| + \gamma)``.  The output of
    the returned function is the sum of the scaled per-example outputs.

    Args:
        fun: The per-example function to be transformed.
        has_aux: If True, ``fun`` returns ``(value, aux)``; only ``value``
            is scaled and aggregated.
        batch_argnums: Which arguments have a batch dimension.
        R: Sensitivity max_norm for the scaled output.  When ``PerGroup``,
            each group is scaled independently.
        gamma: Denominator stabilizer :math:`\gamma` (default 0.01).
        normalize_by: Divisor applied to the scaled sum.
        return_aux: If True, the returned callable returns an
            :class:`AutoClippedFunAux` alongside the summed value.
        return_stats: If True, return :class:`ClippingStats` alongside the
            summed value. Cannot be combined with ``return_aux``.
        microbatch_size: Process the batch in chunks of this size.
        dtype: Optional accumulation dtype for the sum.

    Returns:
        ``(auto_fn, state)``.  ``auto_fn`` has the signature
        ``(*args, state, **kwargs) -> (value, state)`` or
        ``(*args, state, **kwargs) -> ((value, aux), state)`` when
        ``return_aux=True``.

    Raises:
        ConfigurationError: If the batch is empty.

    Formal guarantee:
        Under add/remove or zero-out DP, the L2 sensitivity of the first
        output with respect to the batch arguments is the returned
        ``ClippedPytree.max_norm`` metadata.  The bound is constant and
        data-independent, so per-step Gaussian calibration is correct.
        ``mf_gaussian_noise`` additionally requires that ``max_norm`` not drift across
        steps; AUTO-S satisfies that latch because ``R`` and ``normalize_by``
        are fixed.
    """
    _validate_auto_params(R, gamma)

    scale_fn = _make_auto_scale_fn(R, gamma)
    builtin_scale = _BuiltinScale(kind="auto_s", gamma=gamma)
    inner_fn, _ = clipped_fun(
        fun,
        has_aux=has_aux,
        batch_argnums=batch_argnums,
        clipping_norm=R,
        normalize_by=normalize_by,
        return_aux=return_aux,
        return_stats=return_stats,
        microbatch_size=microbatch_size,
        dtype=dtype,
        _scale_fn=scale_fn,
        _builtin_scale=builtin_scale,
        _chunk_compiler=_chunk_compiler,
    )

    state = AutoClipState()

    if not return_aux:

        def auto_fn(*args, state, **kwargs):
            result, _ = inner_fn(*args, state=None, **kwargs)
            return result, state

        return auto_fn, state

    def auto_fn(*args, state, **kwargs):
        (result, fun_aux), _ = inner_fn(*args, state=None, **kwargs)
        aux = AutoClippedFunAux(
            values=fun_aux.values,
            norms=fun_aux.norms,
            clipped_norms=fun_aux.clipped_norms,
            value_aux=fun_aux.value_aux,
            clipping_rate=fun_aux.clipping_rate,
            batch_size=fun_aux.batch_size,
            group_norms=fun_aux.group_norms,
        )
        return (result, aux), state

    return auto_fn, state

`auto_scale_pytree(pytree, R=1.0, gamma=0.01, *, compute_dtype=None)` ¶

AUTO-S automatic scaling of a PyTree (Bu et al., NeurIPS 2023).

Scales the PyTree by R / (\|pytree\| + gamma) so the output L2 norm is clipped by R for any input. Unlike :func:clip_pytree, there is no threshold to tune — every example contributes an approximately unit-length update, with the effective step size absorbed into the learning rate.

Parameters:

Name	Type	Description	Default
`pytree`	`Any`	Tensor pytree to scale (flat or nested).	required
`R`	`float \| PerGroup`	Output sensitivity bound (non-negative). When `PerGroup`, each group is scaled independently to its own bound; leaf paths must match `R.groups` exactly.	`1.0`
`gamma`	`float`	Small positive denominator stabilizer :math:`\gamma` (default 0.01). Must be strictly positive; at `gamma=0` this reduces to AUTO-V (undefined at zero gradient).	`0.01`
`compute_dtype`	`dtype \| None`	Internal dtype for the L2-norm reduction and leaf scaling. `None` (default) auto-promotes bf16/fp16 inputs to float32.	`None`

Returns:

Type	Description
`Any`	Tuple of (scaled_pytree, aux) where `aux.norm` is the original L2
`ClipPytreeAux`	norm and `aux.group_norms` carries per-group norms in per-group
`tuple[Any, ClipPytreeAux]`	mode.

Formal guarantee

For every input, the output has L2 norm at most R (scalar case) or :math:\sqrt{\sum_k R_k^2} (per-group case). The bound holds on the values as stored: each leaf is scaled in compute_dtype (never below its storage precision) and cast once to its output dtype. The scale is shrunk by the final cast's rounding margin so the result cannot exceed R. Unlike :func:clip_pytree there is no min(1, ...), so the shrink applies to every input — roughly 0.4% of the output magnitude at bfloat16, 0.05% at float16, and a few ULPs at float32.

NaN/Inf values are sanitized to zero before scaling, matching the behavior of :func:clip_pytree.

Source code in packages/opake-engine/src/opake/api/engine/clipping/_pytree.py

def auto_scale_pytree(
    pytree: Any,
    R: float | PerGroup = 1.0,
    gamma: float = 0.01,
    *,
    compute_dtype: torch.dtype | None = None,
) -> tuple[Any, ClipPytreeAux]:
    r"""AUTO-S automatic scaling of a PyTree (Bu et al., NeurIPS 2023).

    Scales the PyTree by ``R / (\|pytree\| + gamma)`` so the output L2
    norm is clipped by ``R`` for any input. Unlike :func:`clip_pytree`, there
    is no threshold to tune — every example contributes an approximately
    unit-length update, with the effective step size absorbed into the
    learning rate.

    Args:
        pytree: Tensor pytree to scale (flat or nested).
        R: Output sensitivity bound (non-negative). When ``PerGroup``, each
            group is scaled independently to its own bound; leaf paths must
            match ``R.groups`` exactly.
        gamma: Small positive denominator stabilizer :math:`\gamma` (default
            0.01). Must be strictly positive; at ``gamma=0`` this reduces to
            AUTO-V (undefined at zero gradient).
        compute_dtype: Internal dtype for the L2-norm reduction and leaf
            scaling. ``None`` (default) auto-promotes bf16/fp16 inputs to
            float32.

    Returns:
        Tuple of (scaled_pytree, aux) where ``aux.norm`` is the original L2
        norm and ``aux.group_norms`` carries per-group norms in per-group
        mode.

    Formal guarantee:
        For every input, the output has L2 norm at most ``R`` (scalar case)
        or :math:`\sqrt{\sum_k R_k^2}` (per-group case).  The bound holds on
        the values as *stored*: each leaf is scaled in ``compute_dtype`` (never
        below its storage precision) and cast once to its output dtype.  The
        scale is shrunk by the final cast's rounding margin so the result cannot
        exceed ``R``. Unlike :func:`clip_pytree` there is no ``min(1, ...)``, so
        the shrink applies to every input — roughly 0.4% of the output magnitude
        at bfloat16, 0.05% at float16, and a few ULPs at float32.

    NaN/Inf values are sanitized to zero before scaling, matching the
    behavior of :func:`clip_pytree`.
    """
    if gamma <= 0:
        raise ConfigurationError(*(f"gamma must be positive, got {gamma}",))

    pytree = tree_map(
        lambda t: (
            torch.nan_to_num(t, nan=0.0, posinf=0.0, neginf=0.0)
            if isinstance(t, torch.Tensor)
            else t
        ),
        pytree,
    )

    if isinstance(R, PerGroup):
        return _auto_scale_per_group(pytree, R, gamma, compute_dtype)

    leaves = [leaf for leaf in tree_leaves(pytree) if isinstance(leaf, torch.Tensor)]
    acc_dtype = _resolve_compute_dtype_for_reduction(leaves, compute_dtype)
    sq_dtype = _sq_accum_dtype(leaves)
    roundoff = _norm_roundoff(
        acc_dtype, sq_dtype, len(leaves), _reduction_terms(leaves)
    )

    norm = torch.sqrt(_sq_norm(leaves, acc_dtype, sq_dtype))
    orig_norm = norm.to(acc_dtype)

    R_tensor = torch.clamp(
        torch.as_tensor(R, dtype=norm.dtype, device=norm.device), min=0.0
    )
    gamma_tensor = torch.tensor(gamma, dtype=norm.dtype, device=norm.device)
    ratio = R_tensor / (norm + gamma_tensor)

    def scale_leaf(t):
        if not isinstance(t, torch.Tensor):
            return t
        return _scale_tensor(t, ratio, acc_dtype, roundoff, clamp_to_one=False)

    scaled = tree_map(
        lambda t: scale_leaf(t) if isinstance(t, torch.Tensor) else t, pytree
    )

    return scaled, ClipPytreeAux(norm=orig_norm, group_norms=None)

`clip_pytree(pytree, clipping_norm, return_zero=False, *, compute_dtype=None)` ¶

Clip a PyTree of tensors to a maximum L2 norm.

NaN and Inf values in the input are replaced with zeros before clipping. This is vmap-compatible and DP-safe: norm(output) <= clipping_norm holds on the values as stored, not just in exact arithmetic. The norm reduction, the multiplication run in compute_dtype (never below the leaf's storage precision), and the result is cast once to the leaf dtype. That cast can round up, so each leaf's scale is shrunk by its final-rounding margin. The shrink is applied before the min(1, ...), so an input well inside the threshold is returned unchanged; one whose norm sits within a few ULPs of clipping_norm is scaled by that much. For clipped inputs the cost is roughly 0.4% of the output magnitude at bfloat16, 0.05% at float16, and a few ULPs at float32.

Parameters:

Name	Type	Description	Default
`pytree`	`Any`	Tensor pytree to clip (flat or nested).	required
`clipping_norm`	`float \| PerGroup`	Maximum L2 norm (non-negative, or inf for no clipping). When `PerGroup`, each group is clipped independently; leaf paths must match `clipping_norm.groups` exactly.	required
`return_zero`	`bool`	If True, the output PyTree is guaranteed to be zero no matter what the inputs are. Does not influence the formal guarantees but useful for privacy amplification via padding (see https://arxiv.org/pdf/2411.04205).	`False`
`compute_dtype`	`dtype \| None`	Internal dtype for the L2-norm reduction and leaf scaling. `None` (default) auto-promotes bf16/fp16 inputs to float32. This keeps the sensitivity bound numerically honest under low-precision compute — small per-element contributions don't get rounded away during the sum-of-squares or multiplication.	`None`

Returns:

Type	Description
`tuple[Any, ClipPytreeAux]`	Tuple of (clipped_pytree, aux) where aux contains: - norm: The L2 norm of the original (unclipped) pytree

Edge cases

clipping_norm=0: Returns zeros
clipping_norm=inf: No clipping (passthrough)
pytree_norm=0: Returns unchanged
NaN/Inf values: Replaced with zeros before clipping
return_zero=True: Returns zeros regardless of other parameters

Source code in packages/opake-engine/src/opake/api/engine/clipping/_pytree.py

def clip_pytree(
    pytree: Any,
    clipping_norm: float | PerGroup,
    return_zero: bool = False,
    *,
    compute_dtype: torch.dtype | None = None,
) -> tuple[Any, ClipPytreeAux]:
    """Clip a PyTree of tensors to a maximum L2 norm.

    NaN and Inf values in the input are replaced with zeros before clipping.
    This is vmap-compatible and DP-safe: ``norm(output) <= clipping_norm`` holds
    on the values as *stored*, not just in exact arithmetic.  The norm reduction,
    the multiplication run in ``compute_dtype`` (never below the leaf's storage
    precision), and the result is cast once to the leaf dtype. That cast can
    round up, so each leaf's scale is shrunk by its final-rounding margin. The
    shrink is applied before the ``min(1, ...)``, so an input well inside the
    threshold is returned unchanged; one whose norm sits within a few ULPs of
    ``clipping_norm`` is scaled by that much. For clipped inputs the cost is
    roughly 0.4% of the output magnitude at bfloat16, 0.05% at float16, and a
    few ULPs at float32.

    Args:
        pytree: Tensor pytree to clip (flat or nested).
        clipping_norm: Maximum L2 norm (non-negative, or inf for no clipping).
            When ``PerGroup``, each group is clipped independently; leaf paths
            must match ``clipping_norm.groups`` exactly.
        return_zero: If True, the output PyTree is guaranteed to be zero no matter
            what the inputs are. Does not influence the formal guarantees but useful
            for privacy amplification via padding (see https://arxiv.org/pdf/2411.04205).
        compute_dtype: Internal dtype for the L2-norm reduction and leaf
            scaling. ``None`` (default) auto-promotes bf16/fp16 inputs to
            float32. This keeps the sensitivity bound numerically honest under
            low-precision compute — small per-element contributions don't get
            rounded away during the sum-of-squares or multiplication.

    Returns:
        Tuple of (clipped_pytree, aux) where aux contains:
            - norm: The L2 norm of the original (unclipped) pytree

    Edge cases:
        - clipping_norm=0: Returns zeros
        - clipping_norm=inf: No clipping (passthrough)
        - pytree_norm=0: Returns unchanged
        - NaN/Inf values: Replaced with zeros before clipping
        - return_zero=True: Returns zeros regardless of other parameters
    """
    pytree = tree_map(
        lambda t: (
            torch.nan_to_num(t, nan=0.0, posinf=0.0, neginf=0.0)
            if isinstance(t, torch.Tensor)
            else t
        ),
        pytree,
    )

    if isinstance(clipping_norm, PerGroup):
        return _clip_pytree_per_group(pytree, clipping_norm, return_zero, compute_dtype)

    leaves = [leaf for leaf in tree_leaves(pytree) if isinstance(leaf, torch.Tensor)]
    acc_dtype = _resolve_compute_dtype_for_reduction(leaves, compute_dtype)
    sq_dtype = _sq_accum_dtype(leaves)
    roundoff = _norm_roundoff(
        acc_dtype, sq_dtype, len(leaves), _reduction_terms(leaves)
    )

    sq_norm = _sq_norm(leaves, acc_dtype, sq_dtype)
    norm = torch.sqrt(sq_norm)
    orig_norm = norm.to(acc_dtype)

    clipping_norm_tensor = torch.as_tensor(
        clipping_norm, dtype=norm.dtype, device=norm.device
    )
    clipping_norm_tensor = torch.clamp(clipping_norm_tensor, min=0.0)
    ratio = clipping_norm_tensor / norm

    def scale_leaf(t):
        if not isinstance(t, torch.Tensor):
            return t
        return _scale_tensor(t, ratio, acc_dtype, roundoff, clamp_to_one=True)

    clipped = tree_map(
        lambda t: scale_leaf(t) if isinstance(t, torch.Tensor) else t, pytree
    )

    if return_zero:
        clipped = tree_map(
            lambda t: torch.zeros_like(t) if isinstance(t, torch.Tensor) else t, clipped
        )

    return clipped, ClipPytreeAux(norm=orig_norm, group_norms=None)

`clipped_fun(fun, has_aux=False, *, batch_argnums=0, clipping_norm=1.0, normalize_by=1.0, return_aux=False, return_stats=False, second_moment=False, microbatch_size=None, dtype=None, compute_dtype=None, _scale_fn=None, _builtin_scale=None, _chunk_compiler=None)` ¶

Transform a function to clip its output and sum across a batch.

This is the primary API for per-example clipping in DP-SGD. It wraps a function to clip each per-example output to a maximum L2 norm, then sums the clipped outputs.

The returned pytree is wrapped as :class:ClippedPytree (single-stream) or :class:SecondMomentClippingOutput (paired-stream when second_moment=True), carrying the post-normalization max_norm for downstream noise calibration. The bound is part of the contract: consumers (gaussian_noise, mf_gaussian_noise) read it directly without the caller threading a separate sensitivity argument. Unwrap to a raw pytree via .pytree if you need the summed values without metadata.

Example Usage

from opake.api.engine.clipping._clipped_fun import clipped_fun data = torch.tensor([0.0, 1.0, 2.0, 3.0, 4.0, 5.0]) clipped_mean, clip_state = clipped_fun(torch.mean, clipping_norm=1.0) result, clip_state = clipped_mean(data, state=clip_state) result.pytree tensor(5.) result.max_norm 1.0

Formal Guarantees

For the first function output: The L2 sensitivity of the returned function with respect to the batch arguments (specified by batch_argnums) under add/remove or zero-out differential privacy definitions is guaranteed to be clipping_norm. Under replace-one DP, the sensitivity is doubled (2 * clipping_norm). Extra auxiliary outputs (aux, norms) are per-example. This function guarantees that per-example outputs only depend on the data for the same example. This allows maximum flexibility for the caller to aggregate these as desired (possibly with a DP mean, median, quantile, or histogram mechanism).

Parameters:

Name	Type	Description	Default
`fun`	`Callable[..., Any]`	The function to be clipped.	required
`has_aux`	`bool`	If True, `fun` is expected to return a tuple `(value, loss_aux)`. Only the value will be clipped + aggregated, `loss_aux` will be returned on a per-example basis. Exercise caution when using this as the sensitivity guarantees of the returned Callable are only provided w.r.t. `value`.	`False`
`batch_argnums`	`int \| tuple[int, ...]`	Specifies which argument(s) of `fun` contain the batch dimension. All arguments specified here must have the same size along the 0th axis.	`0`
`clipping_norm`	`float \| PerGroup`	The maximum L2 norm allowed.	`1.0`
`normalize_by`	`float`	Divide the clipped output by this value before returning.	`1.0`
`return_aux`	`bool`	If True, the returned Callable will return a per-example aux dataclass containing the original per-example values, per-example norms before clipping, and any auxiliary data returned by `fun`.	`False`
`return_stats`	`bool`	If True, return :class:`ClippingStats` with aggregate clipping counts, rates, and batch size. Cannot be combined with `return_aux`.	`False`
`second_moment`	`bool`	If True, also accumulate the element-wise sum of per-example squared clipped values, i.e. `Σᵢ gᵢ²`. The squaring happens inside the per-example loop so the second-stream sensitivity is the per-record squared bound `C²` (averaged: `C² / normalize_by`). The wrapped output becomes :class:`SecondMomentClippingOutput` with both streams. Per-group `clipping_norm` is supported and yields `SecondMomentClippingOutput` with per-group `max_norm` on both streams.	`False`
`microbatch_size`	`int \| None`	If set, the batch is split up into microbatches of this size for memory-efficient processing. Processes each microbatch separately and accumulates results without materializing the full batch of gradients. Set this to reduce peak memory usage at the cost of slightly slower computation. The running sum is held in the wider of `compute_dtype` and the output dtype, so a bf16/fp16 run keeps one float32 copy of the summed output. Pass `compute_dtype=torch.bfloat16` to give that memory back, at the cost of accumulation precision when the output dtype is also low precision.	`None`
`dtype`	`dtype \| None`	Optional dtype for the clipped+aggregated pytree. If None, the dtype will be the same as the dtypes of the function output.	`None`
`compute_dtype`	`dtype \| None`	Internal reduction and accumulation dtype. `None` (default) auto-promotes bf16/fp16 to float32 for numerical stability; an explicit dtype selects the reduction precision. Leaf scaling and microbatch accumulation use the wider of this dtype and the leaf or output storage dtype, so explicit lower precision never narrows a wider value. Independent of `dtype` (which controls the output dtype). Applies across microbatches too, so microbatched and non-microbatched runs agree at their resolved accumulation precision.	`None`

Returns: A tuple (clip_fn, FixedClipState) where clip_fn(*args, state=...) clips the output of fun and sums across the batch. The exact return shape depends on second_moment and return_aux:

| ``second_moment`` | ``return_aux`` | ``clip_fn`` returns                                |
| :---------------- | :------------- | :------------------------------------------------- |
| False             | False          | ``(ClippedPytree, state)``                         |
| False             | True           | ``((ClippedPytree, ClippedFunAux), state)``        |
| True              | False          | ``(SecondMomentClippingOutput, state)``            |
| True              | True           | ``((SecondMomentClippingOutput, ClippedFunAux), state)`` |

With ``return_stats=True`` and ``return_aux=False``, the clipped value
is paired with :class:`ClippingStats`.

Raises:

Type	Description
`ConfigurationError`	If the batch is empty or a chunk compiler is used without a finite microbatch size.

Source code in packages/opake-engine/src/opake/api/engine/clipping/_clipped_fun.py

def clipped_fun(
    fun: Callable[..., Any],
    has_aux: bool = False,
    *,
    batch_argnums: int | tuple[int, ...] = 0,
    clipping_norm: float | PerGroup = 1.0,
    normalize_by: float = 1.0,
    return_aux: bool = False,
    return_stats: bool = False,
    second_moment: bool = False,
    microbatch_size: int | None = None,
    dtype: torch.dtype | None = None,
    compute_dtype: torch.dtype | None = None,
    _scale_fn: Callable | None = None,
    _builtin_scale: _BuiltinScale | None = None,
    _chunk_compiler: Callable | None = None,
) -> tuple[Callable, FixedClipState]:
    """Transform a function to clip its output and sum across a batch.

    This is the primary API for per-example clipping in DP-SGD. It wraps a function
    to clip each per-example output to a maximum L2 norm, then sums the clipped outputs.

    The returned pytree is wrapped as :class:`ClippedPytree` (single-stream)
    or :class:`SecondMomentClippingOutput` (paired-stream when
    ``second_moment=True``), carrying the post-normalization
    ``max_norm`` for downstream noise calibration.  The bound is part of
    the contract: consumers (``gaussian_noise``, ``mf_gaussian_noise``) read it
    directly without the caller threading a separate ``sensitivity``
    argument.  Unwrap to a raw pytree via ``.pytree`` if you need the
    summed values without metadata.

    Example Usage:
        >>> from opake.api.engine.clipping._clipped_fun import clipped_fun
        >>> data = torch.tensor([0.0, 1.0, 2.0, 3.0, 4.0, 5.0])
        >>> clipped_mean, clip_state = clipped_fun(torch.mean, clipping_norm=1.0)
        >>> result, clip_state = clipped_mean(data, state=clip_state)
        >>> result.pytree
        tensor(5.)
        >>> result.max_norm
        1.0

    Formal Guarantees:
        For the first function output:
          The L2 sensitivity of the returned function with respect to the batch
          arguments (specified by `batch_argnums`) under add/remove or zero-out
          differential privacy definitions is guaranteed to be `clipping_norm`.
          Under replace-one DP, the sensitivity is doubled (2 * `clipping_norm`).
        Extra auxiliary outputs (aux, norms) are per-example. This function
          guarantees that per-example outputs only depend on the data for the same
          example. This allows maximum flexibility for the caller to aggregate
          these as desired (possibly with a DP mean, median, quantile, or histogram
          mechanism).

    Args:
        fun: The function to be clipped.
        has_aux: If True, `fun` is expected to return a tuple `(value, loss_aux)`. Only
            the value will be clipped + aggregated, `loss_aux` will be returned on a
            per-example basis. Exercise caution when using this as the sensitivity
            guarantees of the returned Callable are only provided w.r.t. `value`.
        batch_argnums: Specifies which argument(s) of `fun` contain the batch
            dimension. All arguments specified here must have the same size along the
            0th axis.
        clipping_norm: The maximum L2 norm allowed.
        normalize_by: Divide the clipped output by this value before returning.
        return_aux: If True, the returned Callable will return a per-example aux
            dataclass containing the original per-example values, per-example norms
            before clipping, and any auxiliary data returned by `fun`.
        return_stats: If True, return :class:`ClippingStats` with aggregate
            clipping counts, rates, and batch size. Cannot be combined with
            ``return_aux``.
        second_moment: If True, also accumulate the element-wise sum of
            per-example squared clipped values, i.e. ``Σᵢ gᵢ²``.  The
            squaring happens inside the per-example loop so the
            second-stream sensitivity is the per-record squared bound
            ``C²`` (averaged: ``C² / normalize_by``).  The wrapped
            output becomes :class:`SecondMomentClippingOutput` with both
            streams.  Per-group ``clipping_norm`` is supported and yields
            ``SecondMomentClippingOutput`` with per-group ``max_norm``
            on both streams.
        microbatch_size: If set, the batch is split up into microbatches of this
            size for memory-efficient processing. Processes each microbatch separately
            and accumulates results without materializing the full batch of gradients.
            Set this to reduce peak memory usage at the cost of slightly slower computation.
            The running sum is held in the wider of ``compute_dtype`` and the
            output dtype, so a bf16/fp16 run keeps one float32 copy of the summed
            output. Pass ``compute_dtype=torch.bfloat16`` to give that memory back,
            at the cost of accumulation precision when the output dtype is also low
            precision.
        dtype: Optional dtype for the clipped+aggregated pytree. If None, the dtype
            will be the same as the dtypes of the function output.
        compute_dtype: Internal reduction and accumulation dtype. ``None``
            (default) auto-promotes bf16/fp16 to float32 for numerical stability;
            an explicit dtype selects the reduction precision. Leaf scaling and
            microbatch accumulation use the wider of this dtype and the leaf or
            output storage dtype, so explicit lower precision never narrows a wider
            value. Independent of ``dtype`` (which controls the *output* dtype).
            Applies across microbatches too, so microbatched and non-microbatched
            runs agree at their resolved accumulation precision.
    Returns:
        A tuple ``(clip_fn, FixedClipState)`` where ``clip_fn(*args, state=...)``
        clips the output of ``fun`` and sums across the batch.  The exact
        return shape depends on ``second_moment`` and ``return_aux``:

        | ``second_moment`` | ``return_aux`` | ``clip_fn`` returns                                |
        | :---------------- | :------------- | :------------------------------------------------- |
        | False             | False          | ``(ClippedPytree, state)``                         |
        | False             | True           | ``((ClippedPytree, ClippedFunAux), state)``        |
        | True              | False          | ``(SecondMomentClippingOutput, state)``            |
        | True              | True           | ``((SecondMomentClippingOutput, ClippedFunAux), state)`` |

        With ``return_stats=True`` and ``return_aux=False``, the clipped value
        is paired with :class:`ClippingStats`.

    Raises:
        ConfigurationError: If the batch is empty or a chunk compiler is used
            without a finite microbatch size.
    """
    batch_argnums = _prepare_clipped_fun(
        batch_argnums, clipping_norm, return_aux, return_stats
    )
    stream_impl = _resolve_stream_impl(
        clipping_norm=clipping_norm,
        second_moment=second_moment,
        compute_dtype=compute_dtype,
        scale_fn=_scale_fn,
        chunk_compiler=_chunk_compiler,
    )

    # Wrap function to handle has_aux - use empty tuple () not None!
    if not has_aux:

        def fun_with_aux(*args, **kwargs):
            return (fun(*args, **kwargs), ())

    else:
        fun_with_aux = fun

    clip_state = FixedClipState()

    def _per_example_fn(kernel_clipping_norm, *args_single, **call_kwargs):
        value, aux = fun_with_aux(*args_single, **call_kwargs)
        if _scale_fn is None:
            clipped_value, norm = clip_pytree(
                value,
                clipping_norm=kernel_clipping_norm,
                compute_dtype=compute_dtype,
            )
        else:
            clipped_value, norm = _scale_fn(value)
        squared_value = (
            tree_map(
                lambda x: x.square() if isinstance(x, torch.Tensor) else x,
                clipped_value,
            )
            if second_moment
            else None
        )
        if return_aux or return_stats:
            diagnostics = {"norms": norm.norm.detach()}
            if return_aux:
                diagnostics["clipped_norms"] = global_norm(
                    clipped_value, compute_dtype=compute_dtype
                ).detach()
            if norm.group_norms is not None:
                diagnostics["group_norms"] = {
                    key: group_norm.detach()
                    for key, group_norm in norm.group_norms.items()
                }

            if return_aux:
                _attach_value_aux(diagnostics, value, aux, has_aux)

            if second_moment:
                return clipped_value, squared_value, diagnostics
            return clipped_value, diagnostics
        if second_moment:
            return clipped_value, squared_value
        return clipped_value

    if _chunk_compiler is not None and microbatch_size is None:
        raise ConfigurationError(
            *("_chunk_compiler requires a finite microbatch_size",)
        )

    def _streaming_kernel(in_dims, kernel_clipping_norm, args, kwargs, reduce_leaf):
        values, value_aux = _vmap_values(
            fun_with_aux, in_dims, args, kwargs, return_aux
        )
        reduced, markers, squared_reduced, squared_markers, diagnostics = (
            _stream_clip_and_sum if stream_impl is None else stream_impl
        )(
            values,
            kernel_clipping_norm,
            scale=_FIXED_SCALE if _scale_fn is None else _builtin_scale,
            batch_size=batch_size_from_args(args, batch_argnums),
            reduce_leaf=reduce_leaf,
            compute_dtype=compute_dtype,
            second_moment=second_moment,
            return_aux=return_aux,
            return_stats=return_stats,
        )
        if return_aux:
            _attach_value_aux(diagnostics, values, value_aux, has_aux)
        return reduced, markers, squared_reduced, squared_markers, diagnostics

    def _make_chunk_kernel(in_dims, stream):
        # Fixed and AUTO-S scaling are construction-time constants.  A
        # stateful scaling parameter must cross this private boundary as a
        # tensor input rather than be captured as changing Python state.
        def _chunk_kernel(kernel_clipping_norm, *chunk_args, **call_kwargs):
            """Run vmap, clipping, diagnostics, and reduction for one chunk."""
            if stream:
                return _streaming_kernel(
                    in_dims,
                    kernel_clipping_norm,
                    chunk_args,
                    call_kwargs,
                    lambda x: torch.sum(
                        x, dim=0, dtype=_accumulation_dtype(x, dtype, compute_dtype)
                    ),
                )
            n_outputs = 1 + int(bool(second_moment)) + int(return_aux or return_stats)
            out_dims = 0 if n_outputs == 1 else (0,) * n_outputs
            vmapped = _vmap(
                _per_example_fn,
                in_dims=(None, *in_dims),
                out_dims=out_dims,
                randomness="same",
            )
            outputs = vmapped(kernel_clipping_norm, *chunk_args, **call_kwargs)
            if n_outputs == 1:
                clipped_values = outputs
                squared_values = None
                diagnostics = ()
            else:
                output_index = 0
                clipped_values = outputs[output_index]
                output_index += 1
                squared_values = outputs[output_index] if second_moment else None
                if second_moment:
                    output_index += 1
                diagnostics = (
                    outputs[output_index] if (return_aux or return_stats) else ()
                )

            reduced = tree_map(
                lambda x: torch.sum(
                    x,
                    dim=0,
                    dtype=_accumulation_dtype(x, dtype, compute_dtype),
                ),
                clipped_values,
            )
            dtype_markers = tree_map(lambda x: x.new_zeros(()), clipped_values)
            if second_moment:
                squared_reduced = tree_map(
                    lambda x: torch.sum(
                        x,
                        dim=0,
                        dtype=_accumulation_dtype(x, dtype, compute_dtype),
                    ),
                    squared_values,
                )
                squared_dtype_markers = tree_map(
                    lambda x: x.new_zeros(()), squared_values
                )
            else:
                squared_reduced = ()
                squared_dtype_markers = ()
            return (
                reduced,
                dtype_markers,
                squared_reduced,
                squared_dtype_markers,
                diagnostics,
            )

        return _chunk_kernel

    # Cache by positional batching structure so callers may use any valid
    # positional/default-argument arrangement without reusing incompatible
    # vmap ``in_dims``.
    chunk_kernels: dict[tuple[tuple[int | None, ...], bool], Callable] = {}

    def clipped_fn(runtime_clipping_norm, *args, **kwargs):
        current_clipping_norm, output_max_norm, output_squared_max_norm = (
            _resolve_runtime_clipping_norm(
                clipping_norm,
                runtime_clipping_norm,
                normalize_by,
                second_moment,
            )
        )
        kernel_clipping_norm = _prepare_kernel_clipping_norm(
            current_clipping_norm, args, batch_argnums
        )
        in_dims = tuple(0 if i in batch_argnums else None for i in range(len(args)))
        stream = not under_differentiating_transform(
            when_compiling=True
        ) and _streaming_supported(current_clipping_norm, _scale_fn, _builtin_scale)

        # Choose execution path based on microbatch_size
        stats = None
        if microbatch_size is None and stream:
            result, _, squared_result, _, aux = _streaming_kernel(
                in_dims,
                kernel_clipping_norm,
                args,
                kwargs,
                lambda x: _sum_clipped_tensor(
                    x, dim=0, output_dtype=dtype, compute_dtype=compute_dtype
                ),
            )
        elif microbatch_size is None:
            # Fast path: vmap entire batch at once.  Output shape depends
            # on the (second_moment, return_aux) flags — see the per_example_fn
            # branches above.  When n_outputs == 1, vmap returns the single
            # pytree (which may itself be a tuple of tensors for tuple
            # params); when n_outputs > 1 the per_example_fn returns a
            # tuple of n_outputs pytrees.
            n_outputs = 1 + int(bool(second_moment)) + int(return_aux or return_stats)
            out_dims = 0 if n_outputs == 1 else (0,) * n_outputs
            vmapped = _vmap(
                _per_example_fn,
                in_dims=(None, *in_dims),
                out_dims=out_dims,
                randomness="same",
            )
            outputs = vmapped(kernel_clipping_norm, *args, **kwargs)
            if n_outputs == 1:
                clipped_values = outputs
                squared_values = None
                aux = ()
            else:
                idx = 0
                clipped_values = outputs[idx]
                idx += 1
                squared_values = outputs[idx] if second_moment else None
                if second_moment:
                    idx += 1
                aux = outputs[idx] if (return_aux or return_stats) else ()

            # Sum clipped values across batch dimension
            result = tree_map(
                lambda x: _sum_clipped_tensor(
                    x, dim=0, output_dtype=dtype, compute_dtype=compute_dtype
                ),
                clipped_values,
            )
            squared_result = (
                tree_map(
                    lambda x: _sum_clipped_tensor(
                        x, dim=0, output_dtype=dtype, compute_dtype=compute_dtype
                    ),
                    squared_values,
                )
                if second_moment
                else None
            )
        else:
            kernel_key = (in_dims, stream)
            chunk_kernel = chunk_kernels.get(kernel_key)
            if chunk_kernel is None:
                eager_chunk_kernel = _make_chunk_kernel(in_dims, stream)
                if _chunk_compiler is None:
                    chunk_kernel = eager_chunk_kernel
                else:
                    # Compilation policy belongs to the owner; the kernel
                    # itself remains tensor-only. Static chunk variants bound
                    # recompilation by ``microbatch_size`` rather than by every
                    # realized Poisson batch size.
                    chunk_kernel = _chunk_compiler(eager_chunk_kernel)
                chunk_kernels[kernel_key] = chunk_kernel
            result, squared_result, aux, stats = _microbatch_accumulate_reduced(
                chunk_fn=chunk_kernel,
                kernel_clipping_norm=kernel_clipping_norm,
                args=args,
                kwargs=kwargs,
                batch_argnums=batch_argnums,
                microbatch_size=microbatch_size,
                return_aux=return_aux,
                return_stats=return_stats,
                dtype=dtype,
                clipping_norm=current_clipping_norm,
                second_moment=second_moment,
            )

        # Normalize
        if normalize_by != 1.0:
            result = tree_map(lambda x: x / normalize_by, result)
            if second_moment:
                squared_result = tree_map(lambda x: x / normalize_by, squared_result)

        if second_moment:
            output = SecondMomentClippingOutput(
                grads=ClippedPytree(pytree=result, max_norm=output_max_norm),
                squared_grads=ClippedPytree(
                    pytree=squared_result, max_norm=output_squared_max_norm
                ),
            )
        else:
            output = clipped(result, max_norm=output_max_norm)

        if not return_aux and not return_stats:
            return output

        if return_aux:
            aux_dict = aux if isinstance(aux, dict) else {}
            norms = aux_dict.get("norms")
            group_norms_dict = aux_dict.get("group_norms")
            stats = _compute_clipping_stats(
                norms,
                clipping_norm=current_clipping_norm,
                group_norms_dict=group_norms_dict,
            )

            aux = ClippedFunAux(
                values=aux_dict.get("values"),
                norms=norms,
                clipped_norms=aux_dict.get("clipped_norms"),
                value_aux=aux_dict.get("value_aux"),
                clipping_rate=(
                    stats.clipping_rate
                    if isinstance(stats.clipping_rate, float)
                    or stats.clipping_rate is None
                    else None
                ),
                batch_size=stats.batch_size,
                group_norms=aux_dict.get("group_norms"),
            )
            return output, aux

        if microbatch_size is not None:
            assert stats is not None
            return output, stats

        aux_dict = aux if isinstance(aux, dict) else {}
        return output, _compute_clipping_stats(
            aux_dict.get("norms"),
            clipping_norm=current_clipping_norm,
            group_norms_dict=aux_dict.get("group_norms"),
        )

    # Wrap function to accept and return state
    def stateful_clipped_fn(*args, state, **kwargs):
        runtime_clipping_norm = (
            state.clipping_norm if isinstance(state, _RuntimeClipState) else None
        )
        result = clipped_fn(runtime_clipping_norm, *args, **kwargs)
        return result, state  # State unchanged for fixed clipping

    # Return wrapped function with state
    return stateful_clipped_fn, clip_state

`types` ¶

DP-SGD clipping types façade.

Re-exports adaptive (DP-SGD) and AUTO-S clipping state / aux dataclasses for type annotations from the impl trees.

`ClippedGradResult = ClippedGradValue | tuple[ClippedGradValue, 'ClippedGradAux'] | tuple[ClippedGradValue, 'ClippingStats']` `module-attribute` ¶

Value returned by :class:ClippedGradFn before its updated state.

`AdaptiveClippedGradAux` `dataclass` ¶

Bases: ClippedGradAux

Diagnostic outputs from adaptive_clipped_grad.

All fields are diagnostic — they reflect pre-noise, pre-aggregation values and must not be fed back into private computation. Use the returned ClippedPytree.max_norm metadata for noise calibration.

Inherits all fields from :class:ClippedGradAux.

Source code in packages/opake-dpsgd/src/opake/api/dpsgd/clipping/_adaptive.py

@dataclass(frozen=True)
class AdaptiveClippedGradAux(ClippedGradAux):
    """Diagnostic outputs from adaptive_clipped_grad.

    All fields are diagnostic — they reflect pre-noise, pre-aggregation
    values and must not be fed back into private computation.  Use the
    returned ``ClippedPytree.max_norm`` metadata for noise calibration.

    Inherits all fields from :class:`ClippedGradAux`.
    """

`AdaptiveClipState` `dataclass` ¶

Bases: ClipState

Immutable state for adaptive gradient clipping.

This state is passed explicitly to the clipping function and returned as part of the output, enabling pure functional composition.

The fields are internal implementation details. Adaptive clipping must carry the current threshold, next-step threshold, step counter, RNG/config, and local clipping counts so distributed synchronization can aggregate counts and recompute the global next threshold.

Source code in packages/opake-dpsgd/src/opake/api/dpsgd/clipping/_adaptive.py

@dataclass(frozen=True)
class AdaptiveClipState(ClipState):
    """Immutable state for adaptive gradient clipping.

    This state is passed explicitly to the clipping function and returned
    as part of the output, enabling pure functional composition.

    The fields are internal implementation details. Adaptive clipping must
    carry the current threshold, next-step threshold, step counter, RNG/config,
    and local clipping counts so distributed synchronization can aggregate
    counts and recompute the global next threshold.
    """

    _current_clipping_norm: float | PerGroup
    _next_clipping_norm: float | PerGroup
    _step: int

    # -- internal (config carried for distributed sync) --
    _rng_key: RngKey
    _fraction_noise_std: float
    _expected_batch_size: float
    _learning_rate: float
    _target_quantile: float
    _clipping_norm_min: float
    _clipping_norm_max: float

    # -- internal (per-step counts for distributed aggregation) --
    # Scalar for global clipping, dict[str, float] for per-group clipping.
    _num_clipped: float | dict[str, float]
    _batch_size: float

    def __post_init__(self):
        """Validate state values."""
        if isinstance(self._next_clipping_norm, PerGroup):
            for gname, val in self._next_clipping_norm.values.items():
                if val <= 0:
                    raise ConfigurationError(
                        *(
                            f"next_clipping_norm must be positive for all groups, "
                            f"got {val} for group '{gname}'",
                        )
                    )
        elif self._next_clipping_norm <= 0:
            raise ConfigurationError(
                *(
                    f"next_clipping_norm must be positive, got {self._next_clipping_norm}",
                )
            )
        if self._expected_batch_size <= 0:
            raise ConfigurationError(
                *(
                    f"expected_batch_size must be positive, got {self._expected_batch_size}",
                )
            )
        if self._fraction_noise_std <= 0:
            raise ConfigurationError(
                *(f"fraction_noise_std must be > 0, got {self._fraction_noise_std}",)
            )

`__post_init__()` ¶

Validate state values.

Source code in packages/opake-dpsgd/src/opake/api/dpsgd/clipping/_adaptive.py

def __post_init__(self):
    """Validate state values."""
    if isinstance(self._next_clipping_norm, PerGroup):
        for gname, val in self._next_clipping_norm.values.items():
            if val <= 0:
                raise ConfigurationError(
                    *(
                        f"next_clipping_norm must be positive for all groups, "
                        f"got {val} for group '{gname}'",
                    )
                )
    elif self._next_clipping_norm <= 0:
        raise ConfigurationError(
            *(
                f"next_clipping_norm must be positive, got {self._next_clipping_norm}",
            )
        )
    if self._expected_batch_size <= 0:
        raise ConfigurationError(
            *(
                f"expected_batch_size must be positive, got {self._expected_batch_size}",
            )
        )
    if self._fraction_noise_std <= 0:
        raise ConfigurationError(
            *(f"fraction_noise_std must be > 0, got {self._fraction_noise_std}",)
        )

`AutoClippedFunAux` `dataclass` ¶

Bases: ClippedFunAux

Diagnostic outputs from auto_clipped_fun.

Inherits all fields from :class:ClippedFunAux with AUTO-S semantics:

norms: per-example L2 norms before scaling (i.e. ||v||).
clipped_norms: per-example L2 norms after AUTO-S scaling (R * ||v|| / (||v|| + gamma)), clipped by R.
clipping_rate: fraction of examples with ||v|| > R (approximately the fraction where the scale factor was < 1).
values, value_aux, batch_size, group_norms: unchanged from :class:ClippedFunAux.

Source code in packages/opake-engine/src/opake/api/engine/clipping/_auto.py

@dataclass(frozen=True)
class AutoClippedFunAux(ClippedFunAux):
    """Diagnostic outputs from ``auto_clipped_fun``.

    Inherits all fields from :class:`ClippedFunAux` with AUTO-S
    semantics:

    - ``norms``: per-example L2 norms before scaling (i.e. ``||v||``).
    - ``clipped_norms``: per-example L2 norms after AUTO-S scaling
      (``R * ||v|| / (||v|| + gamma)``), clipped by ``R``.
    - ``clipping_rate``: fraction of examples with ``||v|| > R``
      (approximately the fraction where the scale factor was < 1).
    - ``values``, ``value_aux``, ``batch_size``, ``group_norms``:
      unchanged from :class:`ClippedFunAux`.
    """

`AutoClippedGradAux` `dataclass` ¶

Bases: ClippedGradAux

Diagnostic outputs from auto_clipped_grad.

Inherits all fields from :class:ClippedGradAux with AUTO-S semantics:

grad_norms: per-example gradient L2 norms before scaling.
clipped_grad_norms: per-example gradient L2 norms after AUTO-S scaling (R * ||g|| / (||g|| + gamma)), clipped by R.
clipping_rate: fraction of examples with ||g|| > R (approximately the fraction where the scale factor was < 1).
loss_values, loss_aux, batch_size, group_norms: unchanged from :class:ClippedGradAux.

Source code in packages/opake-engine/src/opake/api/engine/clipping/_auto.py

@dataclass(frozen=True)
class AutoClippedGradAux(ClippedGradAux):
    """Diagnostic outputs from ``auto_clipped_grad``.

    Inherits all fields from :class:`ClippedGradAux` with AUTO-S
    semantics:

    - ``grad_norms``: per-example gradient L2 norms before scaling.
    - ``clipped_grad_norms``: per-example gradient L2 norms after
      AUTO-S scaling (``R * ||g|| / (||g|| + gamma)``), clipped by ``R``.
    - ``clipping_rate``: fraction of examples with ``||g|| > R``
      (approximately the fraction where the scale factor was < 1).
    - ``loss_values``, ``loss_aux``, ``batch_size``, ``group_norms``:
      unchanged from :class:`ClippedGradAux`.
    """

`AutoClipState` `dataclass` ¶

Bases: ClipState

Marker state for AUTO-S automatic clipping.

Source code in packages/opake-engine/src/opake/api/engine/clipping/_auto.py

@dataclass(frozen=True)
class AutoClipState(ClipState):
    """Marker state for AUTO-S automatic clipping."""

`ClippedGradFn` ¶

Bases: Protocol

Callable returned by :func:opake.dpsgd.clipping.clipped_grad.

The positional arguments mirror the wrapped loss function. state is passed by keyword and every invocation returns a new immutable :class:FixedClipState with the clipped gradient result.

Source code in packages/opake-engine/src/opake/api/engine/clipping/_types.py

class ClippedGradFn(Protocol):
    """Callable returned by :func:`opake.dpsgd.clipping.clipped_grad`.

    The positional arguments mirror the wrapped loss function. ``state`` is
    passed by keyword and every invocation returns a new immutable
    :class:`FixedClipState` with the clipped gradient result.
    """

    def __call__(
        self,
        *args: Any,
        state: FixedClipState,
        **kwargs: Any,
    ) -> tuple[ClippedGradResult, FixedClipState]: ...

`ClippingStats` `dataclass` ¶

Aggregated clipping statistics without per-example materialization.

Source code in packages/opake-engine/src/opake/api/engine/clipping/_clipped_fun.py

@dataclass(frozen=True)
class ClippingStats:
    """Aggregated clipping statistics without per-example materialization."""

    num_clipped: float | dict[str, float]
    clipping_rate: float | dict[str, float] | None
    batch_size: int = 0

`FixedClipState` `dataclass` ¶

Bases: ClipState

Marker state for fixed (non-adaptive) clipping.

Returned by :func:clipped_fun and :func:opake.api.engine.clipping.clipped_grad. Carries no fields; the configured clipping threshold flows through the ClippedPytree.max_norm metadata, not through the state.

Source code in packages/opake-engine/src/opake/api/engine/clipping/_clipped_fun.py

@dataclass(frozen=True)
class FixedClipState(_ClipState):
    """Marker state for fixed (non-adaptive) clipping.

    Returned by :func:`clipped_fun` and :func:`opake.api.engine.clipping.clipped_grad`.
    Carries no fields; the configured clipping threshold flows through
    the ``ClippedPytree.max_norm`` metadata, not through the state.
    """

`clipping` ¶

DP-FTRL clipping façade — fixed threshold, AUTO-S, and per-group norms.

Adaptive thresholding is only available under :mod:opake.dpsgd.clipping (DP-SGD). State and aux types live in :mod:opake.dpftrl.clipping.types.

`auto_clipped_grad(loss_fn, argnums=0, has_aux=False, *, R=1.0, gamma=_DEFAULT_GAMMA, normalize_by=1.0, batch_argnums=1, return_aux=False, return_stats=False, pre_clipping_transform=lambda x: x, microbatch_size=None, dtype=None, second_moment=False, _chunk_compiler=None)` ¶

Create a function that computes the sum of AUTO-S scaled per-example gradients.

AUTO-S (Bu et al., 2023) replaces the standard clipping threshold with a per-example scaling

.. math::

\tilde g_i = R \cdot g_i / (\lVert g_i \rVert + \gamma)

so every per-example gradient has L2 norm at most R by construction. There is no threshold to tune; the effective step size is absorbed into the optimizer learning rate.

Parameters:

Name	Type	Description	Default
`loss_fn`	`Callable`	Scalar loss function. If `has_aux`, returns `(scalar, loss_aux)`.	required
`argnums`	`int \| tuple[int, ...]`	Which argument(s) to differentiate w.r.t.	`0`
`has_aux`	`bool`	If True, `loss_fn` returns `(scalar, loss_aux)`.	`False`
`R`	`float \| PerGroup`	Sensitivity max_norm (default 1.0). When `PerGroup`, each parameter group is scaled independently to its own `R_k`.	`1.0`
`gamma`	`float`	Denominator stabilizer :math:`\gamma` (default 0.01, strictly positive).	`_DEFAULT_GAMMA`
`normalize_by`	`float`	Divisor applied to the summed gradients (set to expected batch size for averaged gradients).	`1.0`
`batch_argnums`	`int \| tuple[int, ...]`	Which arguments have a batch dimension.	`1`
`return_aux`	`bool`	If True, returns per-example diagnostics as :class:`AutoClippedGradAux`.	`False`
`pre_clipping_transform`	`Callable`	An optional function to apply to the per-example gradients before AUTO-S scaling. The function should consume the gradient pytree for a single example and return a new pytree (possibly with different structure). Can be used to e.g., scale the leaves of the pytree to accommodate preconditioner clipping. Does not affect the sensitivity guarantee. Default is identity function.	`lambda x: x`
`microbatch_size`	`int \| None`	Process the batch in chunks of this size.	`None`
`dtype`	`Any`	Optional accumulation dtype for the summed gradient.	`None`
`second_moment`	`bool`	If True, also accumulate the per-example sum of element-wise squared scaled gradients and return a :class:`~opake.types.SecondMomentClippingOutput` carrying both streams. Privacy accounting is unchanged (`gaussian(noise_multiplier)` for DP-SGD or any standard MF mechanism for DP-FTRL); the sensitivity-proportional joint Mahalanobis allocation gives the paired release the same PLD as a single first-moment release.	`False`
`return_stats`	`bool`	If True, return :class:`ClippingStats` beside the gradient. Cannot be combined with `return_aux`.	`False`

Returns:

Type	Description
`Callable`	`(grad_fn, state)`. `grad_fn` has the signature
`AutoClipState`	`(args, state, *kwargs) -> (grad, state)`, or
`tuple[Callable, AutoClipState]`	`(args, state, *kwargs) -> ((grad, aux), state)` when
`tuple[Callable, AutoClipState]`	`return_aux=True`. When `return_stats=True`, it returns
`tuple[Callable, AutoClipState]`	`((grad, stats), state)`.

Formal guarantee

Under add/remove or zero-out DP, the L2 sensitivity of the summed gradients is the returned ClippedPytree.max_norm metadata — a constant independent of the input data. Per-step Gaussian noise (gaussian_noise, …) reads that value each step; PLD composition does not require max_norm to be identical across steps when the accountant models step-varying sensitivity (for example adaptive clipping with adaclip). mf_gaussian_noise's matrix-factorization correlated noise does require max_norm to stay fixed for the whole run — the dispatcher latches the first-call bound — and AUTO-S satisfies that because R and normalize_by do not drift. AUTO-S scaling is per-example and adds no privacy cost beyond gaussian(noise_multiplier) for the gradient release.

Example

import torch from opake.api.engine.clipping import auto_clipped_grad def loss_fn(params, x, y): ... return ((x @ params - y) ** 2).mean() grad_fn, state = auto_clipped_grad( ... loss_fn, argnums=0, batch_argnums=(1, 2), ... R=1.0, normalize_by=32, ... ) params = torch.randn(10) batch_x = torch.randn(32, 10) batch_y = torch.randn(32) grads, state = grad_fn(params, batch_x, batch_y, state=state)

The returned ClippedPytree carries max_norm = R / normalize_by (constant across steps), so it composes with mf_gaussian_noise (DP-FTRL) and gaussian_noise (DP-SGD) the same way clipped_grad does at the same bound.

References

Bu, Wang, Zha, Karypis. "Automatic Clipping: Differentially Private Deep Learning Made Easier and Stronger." NeurIPS 2023. https://arxiv.org/abs/2206.07136

Source code in packages/opake-engine/src/opake/api/engine/clipping/_auto.py

def auto_clipped_grad(
    loss_fn: Callable,
    argnums: int | tuple[int, ...] = 0,
    has_aux: bool = False,
    *,
    R: float | PerGroup = 1.0,
    gamma: float = _DEFAULT_GAMMA,
    normalize_by: float = 1.0,
    batch_argnums: int | tuple[int, ...] = 1,
    return_aux: bool = False,
    return_stats: bool = False,
    pre_clipping_transform: Callable = lambda x: x,
    microbatch_size: int | None = None,
    dtype: Any = None,
    second_moment: bool = False,
    _chunk_compiler: Callable | None = None,
) -> tuple[Callable, AutoClipState]:
    r"""Create a function that computes the sum of AUTO-S scaled per-example gradients.

    AUTO-S (Bu et al., 2023) replaces the standard clipping threshold with
    a per-example scaling

    .. math::

        \tilde g_i = R \cdot g_i / (\lVert g_i \rVert + \gamma)

    so every per-example gradient has L2 norm at most ``R`` by
    construction.  There is no threshold to tune; the effective step size
    is absorbed into the optimizer learning rate.

    Args:
        loss_fn: Scalar loss function.  If ``has_aux``, returns
            ``(scalar, loss_aux)``.
        argnums: Which argument(s) to differentiate w.r.t.
        has_aux: If True, ``loss_fn`` returns ``(scalar, loss_aux)``.
        R: Sensitivity max_norm (default 1.0).  When ``PerGroup``, each
            parameter group is scaled independently to its own ``R_k``.
        gamma: Denominator stabilizer :math:`\gamma` (default 0.01,
            strictly positive).
        normalize_by: Divisor applied to the summed gradients (set to
            expected batch size for averaged gradients).
        batch_argnums: Which arguments have a batch dimension.
        return_aux: If True, returns per-example diagnostics as
            :class:`AutoClippedGradAux`.
        pre_clipping_transform: An optional function to apply to the
            per-example gradients before AUTO-S scaling. The function
            should consume the gradient pytree for a single example
            and return a new pytree (possibly with different structure).
            Can be used to e.g., scale the leaves of the pytree to
            accommodate preconditioner clipping. Does not affect the
            sensitivity guarantee. Default is identity function.
        microbatch_size: Process the batch in chunks of this size.
        dtype: Optional accumulation dtype for the summed gradient.
        second_moment: If True, also accumulate the per-example sum of
            element-wise squared scaled gradients and return a
            :class:`~opake.types.SecondMomentClippingOutput` carrying
            both streams.  Privacy accounting is unchanged
            (``gaussian(noise_multiplier)`` for DP-SGD or any standard MF
            mechanism for DP-FTRL); the sensitivity-proportional joint
            Mahalanobis allocation gives the paired release the same PLD
            as a single first-moment release.
        return_stats: If True, return :class:`ClippingStats` beside the
            gradient. Cannot be combined with ``return_aux``.

    Returns:
        ``(grad_fn, state)``.  ``grad_fn`` has the signature
        ``(*args, state, **kwargs) -> (grad, state)``, or
        ``(*args, state, **kwargs) -> ((grad, aux), state)`` when
        ``return_aux=True``. When ``return_stats=True``, it returns
        ``((grad, stats), state)``.

    Formal guarantee:
        Under add/remove or zero-out DP, the L2 sensitivity of the
        summed gradients is the returned ``ClippedPytree.max_norm`` metadata —
        a constant independent of the input data.  Per-step Gaussian noise
        (``gaussian_noise``, …) reads that value
        each step; PLD composition does not require ``max_norm`` to be
        identical across steps when the accountant models step-varying
        sensitivity (for example adaptive clipping with ``adaclip``).
        ``mf_gaussian_noise``'s matrix-factorization correlated noise *does* require
        ``max_norm`` to stay fixed for the whole run — the dispatcher latches
        the first-call bound — and AUTO-S satisfies that because ``R`` and
        ``normalize_by`` do not drift.  AUTO-S scaling is per-example and adds
        no privacy cost beyond ``gaussian(noise_multiplier)`` for the
        gradient release.

    Example:
        >>> import torch
        >>> from opake.api.engine.clipping import auto_clipped_grad
        >>> def loss_fn(params, x, y):
        ...     return ((x @ params - y) ** 2).mean()
        >>> grad_fn, state = auto_clipped_grad(
        ...     loss_fn, argnums=0, batch_argnums=(1, 2),
        ...     R=1.0, normalize_by=32,
        ... )
        >>> params = torch.randn(10)
        >>> batch_x = torch.randn(32, 10)
        >>> batch_y = torch.randn(32)
        >>> grads, state = grad_fn(params, batch_x, batch_y, state=state)

    The returned ``ClippedPytree`` carries ``max_norm = R / normalize_by``
    (constant across steps), so it composes with ``mf_gaussian_noise`` (DP-FTRL)
    and ``gaussian_noise`` (DP-SGD) the same way ``clipped_grad`` does at
    the same bound.

    References:
        Bu, Wang, Zha, Karypis.  "Automatic Clipping: Differentially
        Private Deep Learning Made Easier and Stronger." NeurIPS 2023.
        https://arxiv.org/abs/2206.07136
    """
    _validate_auto_params(R, gamma)

    scale_fn = _make_auto_scale_fn(R, gamma)
    builtin_scale = _BuiltinScale(kind="auto_s", gamma=gamma)
    inner_fn, _ = clipped_grad(
        loss_fn,
        argnums=argnums,
        has_aux=has_aux,
        clipping_norm=R,
        normalize_by=normalize_by,
        batch_argnums=batch_argnums,
        return_aux=return_aux,
        return_stats=return_stats,
        pre_clipping_transform=pre_clipping_transform,
        microbatch_size=microbatch_size,
        dtype=dtype,
        second_moment=second_moment,
        _scale_fn=scale_fn,
        _builtin_scale=builtin_scale,
        _chunk_compiler=_chunk_compiler,
    )

    state = AutoClipState()

    if not return_aux:

        def grad_fn(*args, state, **kwargs):
            result, _ = inner_fn(*args, state=None, **kwargs)
            return result, state

        return grad_fn, state

    def grad_fn(*args, state, **kwargs):
        inner_result, _ = inner_fn(*args, state=None, **kwargs)
        grads, grad_aux = inner_result
        auto_aux = AutoClippedGradAux(
            loss_values=grad_aux.loss_values,
            grad_norms=grad_aux.grad_norms,
            clipped_grad_norms=grad_aux.clipped_grad_norms,
            loss_aux=grad_aux.loss_aux,
            clipping_rate=grad_aux.clipping_rate,
            batch_size=grad_aux.batch_size,
            group_norms=grad_aux.group_norms,
        )
        return (grads, auto_aux), state

    return grad_fn, state

`clipped_grad(loss_fn, argnums=0, has_aux=False, *, clipping_norm, normalize_by=1.0, batch_argnums=1, return_aux=False, return_stats=False, second_moment=False, pre_clipping_transform=lambda x: x, microbatch_size=None, dtype=None, compute_dtype=None, _scale_fn=None, _builtin_scale=None, _chunk_compiler=None)` ¶

Create a function to compute the sum of clipped gradients of loss_fn.

This function acts as a transformation similar to torch.func.grad, but with added functionality for gradient clipping applied on a per-example (or per-group) basis before summation. It computes the gradient of loss_fn with respect to argnums, calculates the L2 norm of the gradient for each example slice along the first axis of the batch_argnums args, clips each per-example gradient to have a norm of at most clipping_norm, and finally sums these clipped gradients.

Non-grad outputs of the returned function (aux) may optionally be returned by setting return_aux=True. These outputs are per-example, and hence have a batch axis. It is up to the caller to handle these as necessary.

Example Usage

import torch from opake.dpsgd.clipping import clipped_grad f = lambda param, data: 0.5 * ((data - param) ** 2).mean() g, clip_state = clipped_grad(f, clipping_norm=float('inf')) result, clip_state = g(torch.tensor(3.0), torch.tensor([0.0, 7.0, -2.0]), state=clip_state) result tensor(1.3333)

Example Usage (with Auxiliary Output): >>> g, clip_state = clipped_grad( ... f, clipping_norm=float('inf'), return_aux=True ... ) >>> (_, aux), clip_state = g(torch.tensor(3.0), torch.tensor([0.0, 7.0, -2.0]), state=clip_state) >>> aux.loss_values tensor([4.5000, 8.0000, 12.5000]) >>> aux.grad_norms tensor([3., 4., 5.])

Formal Guarantees

For the gradient output: The L2 sensitivity of the returned function with respect to the batch arguments (specified by batch_argnums) under add/remove or zero-out differential privacy definitions is guaranteed to be clipping_norm. Under replace-one DP, the sensitivity is doubled (2 * clipping_norm). All auxiliary outputs (loss_values, grad_norms) are per-example. This function guarantees that per-example outputs only depend on the data for the same example. This allows maximum flexibility for the caller to aggregate these as desired (possibly with a DP mean, median, quantile, or histogram mechanism).

Parameters:

Name	Type	Description	Default
`loss_fn`	`Callable`	The loss function to be differentiated, which should return a scalar value. If `has_aux` is True, it should return a tuple `(value, loss_aux)`.	required
`argnums`	`int \| tuple[int, ...]`	Specifies which argument(s) of `loss_fn` to differentiate with respect to. Can be an integer or a sequence of integers. These arguments should not have a batch dimension.	`0`
`has_aux`	`bool`	If True, `loss_fn` is expected to return a tuple `(value, loss_aux)`. The auxiliary data `loss_aux` will be returned by the transformed function. Exercise caution when using this as no DP sensitivity guarantees are provided for the auxiliary data.	`False`
`clipping_norm`	`float \| PerGroup`	The maximum L2 norm for each per-example gradient. Gradients with a norm larger than this value will be scaled down.	required
`normalize_by`	`float`	Divide the clipped sum by this constant before returning. Set to expected batch size to produce averaged gradients with sensitivity = clipping_norm / normalize_by.	`1.0`
`batch_argnums`	`int \| tuple[int, ...]`	Specifies which argument(s) of `loss_fn` contain the batch dimension (usually the data and labels). Can be an integer or a sequence of integers. All arguments specified here must have the same size along their first dimension (the batch dimension). The default value of 1 assumes the signature of loss_fn is `loss_fn(params, batch)`.	`1`
`return_aux`	`bool`	If True, the transformed function will also return a per-example aux dataclass containing loss values, gradient norms, and loss aux.	`False`
`return_stats`	`bool`	If True, return aggregate clipping statistics including the number and fraction of clipped examples and the batch size. Cannot be combined with `return_aux`.	`False`
`pre_clipping_transform`	`Callable`	An optional function to apply to the per-example gradients before clipping. The function should consume the gradient pytree for a single example and return a new pytree (possibly with different structure). Can be used to e.g., scale the leaves of the pytree to accommodate preconditioner clipping. Account separately for any data-dependent control flow or telemetry. Default is identity function.	`lambda x: x`
`microbatch_size`	`int \| None`	If set, the batch is split up into microbatches of this size for memory-efficient processing. Processes each microbatch separately and accumulates results without materializing the full batch of gradients. Set this to reduce peak memory usage at the cost of slightly slower computation.	`None`
`dtype`	`dtype \| None`	Optional dtype for the returned gradient. If None, the dtype will be the same as the dtypes of the gradient function. Can be useful to avoid overflow issues when using low-precision dtypes as the returned function computes a sum over a potentially large batch.	`None`
`compute_dtype`	`dtype \| None`	Internal reduction and accumulation dtype. `None` (default) auto-promotes bf16/fp16 to float32 for numerical stability; an explicit dtype selects the reduction precision. Leaf scaling and microbatch accumulation use the wider of this dtype and the leaf or output storage dtype, so explicit lower precision never narrows a wider value. Independent of `dtype` (which controls the output dtype).	`None`
`second_moment`	`bool`	Whether to accumulate the clipped-gradient second moment required by DP-FTRL noise mechanisms.	`False`

Returns: Tuple of (:class:ClippedGradFn, clip_state) where: - clipped_grad_fn: A function that computes the sum of clipped per-example gradients. Call signature: clipped_grads, new_state = clipped_grad_fn(..., state=clip_state) If auxiliary outputs are requested, returns: (clipped_grads, grad_aux), new_state If return_stats=True, returns (clipped_grads, stats), new_state. - clip_state: Initial FixedClipState containing sensitivity information

The grad_aux output (when requested) is a ClippedGradAux dataclass with fields:
    - loss_values: Per-example function values (if return_aux=True)
    - grad_norms: Per-example gradient norms (if return_aux=True)
    - loss_aux: Per-example auxiliary data (if has_aux=True)

Source code in packages/opake-engine/src/opake/api/engine/clipping/_clipped_grad.py

def clipped_grad(
    loss_fn: Callable,
    argnums: int | tuple[int, ...] = 0,
    has_aux: bool = False,
    *,
    clipping_norm: float | PerGroup,
    normalize_by: float = 1.0,
    batch_argnums: int | tuple[int, ...] = 1,
    return_aux: bool = False,
    return_stats: bool = False,
    second_moment: bool = False,
    pre_clipping_transform: Callable = lambda x: x,
    microbatch_size: int | None = None,
    dtype: torch.dtype | None = None,
    compute_dtype: torch.dtype | None = None,
    _scale_fn: Callable | None = None,
    _builtin_scale: _BuiltinScale | None = None,
    _chunk_compiler: Callable | None = None,
) -> tuple[ClippedGradFn, FixedClipState]:
    """Create a function to compute the sum of clipped gradients of loss_fn.

    This function acts as a transformation similar to `torch.func.grad`, but with added
    functionality for gradient clipping applied on a per-example (or per-group)
    basis before summation. It computes the gradient of `loss_fn` with respect to
    `argnums`, calculates the L2 norm of the gradient for each example slice
    along the first axis of the `batch_argnums` args, clips each per-example
    gradient to have a norm of at most `clipping_norm`, and finally sums these
    clipped gradients.

    Non-grad outputs of the returned function (aux) may optionally be returned
    by setting `return_aux=True`. These outputs are per-example, and hence have
    a batch axis. It is up to the caller to handle these as necessary.

    Example Usage:
        >>> import torch
        >>> from opake.dpsgd.clipping import clipped_grad
        >>> f = lambda param, data: 0.5 * ((data - param) ** 2).mean()
        >>> g, clip_state = clipped_grad(f, clipping_norm=float('inf'))
        >>> result, clip_state = g(torch.tensor(3.0), torch.tensor([0.0, 7.0, -2.0]), state=clip_state)
        >>> result
        tensor(1.3333)

    Example Usage (with Auxiliary Output):
        >>> g, clip_state = clipped_grad(
        ...     f, clipping_norm=float('inf'), return_aux=True
        ... )
        >>> (_, aux), clip_state = g(torch.tensor(3.0), torch.tensor([0.0, 7.0, -2.0]), state=clip_state)
        >>> aux.loss_values
        tensor([4.5000, 8.0000, 12.5000])
        >>> aux.grad_norms
        tensor([3., 4., 5.])

    Formal Guarantees:
        For the gradient output:
          The L2 sensitivity of the returned function with respect to the batch
          arguments (specified by `batch_argnums`) under add/remove or zero-out
          differential privacy definitions is guaranteed to be `clipping_norm`.
          Under replace-one DP, the sensitivity is doubled (2 * `clipping_norm`).
        All auxiliary outputs (loss_values, grad_norms) are per-example. This
          function guarantees that per-example outputs only depend on the data for the
          same example. This allows maximum flexibility for the caller to aggregate
          these as desired (possibly with a DP mean, median, quantile, or histogram
          mechanism).

    Args:
        loss_fn: The loss function to be differentiated, which should return a scalar
            value. If `has_aux` is True, it should return a tuple `(value, loss_aux)`.
        argnums: Specifies which argument(s) of `loss_fn` to differentiate with respect
            to. Can be an integer or a sequence of integers. These arguments should
            *not* have a batch dimension.
        has_aux: If True, `loss_fn` is expected to return a tuple `(value, loss_aux)`.
            The auxiliary data `loss_aux` will be returned by the transformed function.
            Exercise caution when using this as no DP sensitivity guarantees are
            provided for the auxiliary data.
        clipping_norm: The maximum L2 norm for each per-example gradient. Gradients
            with a norm larger than this value will be scaled down.
        normalize_by: Divide the clipped sum by this constant before returning.
            Set to expected batch size to produce averaged gradients with
            sensitivity = clipping_norm / normalize_by.
        batch_argnums: Specifies which argument(s) of `loss_fn` contain the batch
            dimension (usually the data and labels). Can be an integer or a sequence
            of integers. All arguments specified here must have the same size along
            their first dimension (the batch dimension). The default value of 1 assumes
            the signature of loss_fn is `loss_fn(params, batch)`.
        return_aux: If True, the transformed function will also return a per-example
            aux dataclass containing loss values, gradient norms, and loss aux.
        return_stats: If True, return aggregate clipping statistics including
            the number and fraction of clipped examples and the batch size.
            Cannot be combined with ``return_aux``.
        pre_clipping_transform: An optional function to apply to the per-example
            gradients before clipping. The function should consume the gradient pytree
            for a single example and return a new pytree (possibly with different
            structure). Can be used to e.g., scale the leaves of the pytree to
            accommodate preconditioner clipping. Account separately for any
            data-dependent control flow or telemetry. Default is identity function.
        microbatch_size: If set, the batch is split up into microbatches of this
            size for memory-efficient processing. Processes each microbatch separately
            and accumulates results without materializing the full batch of gradients.
            Set this to reduce peak memory usage at the cost of slightly slower computation.
        dtype: Optional dtype for the returned gradient. If None, the dtype will be
            the same as the dtypes of the gradient function. Can be useful to avoid
            overflow issues when using low-precision dtypes as the returned function
            computes a sum over a potentially large batch.
        compute_dtype: Internal reduction and accumulation dtype. ``None``
            (default) auto-promotes bf16/fp16 to float32 for numerical stability;
            an explicit dtype selects the reduction precision. Leaf scaling and
            microbatch accumulation use the wider of this dtype and the leaf or
            output storage dtype, so explicit lower precision never narrows a wider
            value. Independent of ``dtype`` (which controls the *output* dtype).
        second_moment: Whether to accumulate the clipped-gradient second
            moment required by DP-FTRL noise mechanisms.
    Returns:
        Tuple of (:class:`ClippedGradFn`, clip_state) where:
        - clipped_grad_fn: A function that computes the sum of clipped per-example gradients.
          Call signature: clipped_grads, new_state = clipped_grad_fn(..., state=clip_state)
          If auxiliary outputs are requested, returns: (clipped_grads, grad_aux), new_state
          If ``return_stats=True``, returns ``(clipped_grads, stats), new_state``.
        - clip_state: Initial FixedClipState containing sensitivity information

        The grad_aux output (when requested) is a ClippedGradAux dataclass with fields:
            - loss_values: Per-example function values (if return_aux=True)
            - grad_norms: Per-example gradient norms (if return_aux=True)
            - loss_aux: Per-example auxiliary data (if has_aux=True)
    """
    _validate_static_args(argnums, batch_argnums, normalize_by)
    if return_aux and return_stats:
        raise ConfigurationError(
            *("return_stats cannot be combined with return_aux=True",)
        )

    argnums_tuple = normalize_to_tuple(argnums)
    batch_argnums_tuple = normalize_to_tuple(batch_argnums)
    loss_fn = normalize_fun_to_return_aux(loss_fn, has_aux)

    output_max_norm = clipping_norm / normalize_by
    output_squared_max_norm = (
        (clipping_norm * clipping_norm) / normalize_by if second_moment else None
    )

    def _empty_batch_response(args, state):
        """Return zero gradients and empty diagnostics without vmap."""
        zeros = empty_clipped_grads_like(
            args, argnums_tuple, pre_clipping_transform, dtype
        )
        if second_moment:
            from opake.api.engine.types import SecondMomentClippingOutput

            grads = SecondMomentClippingOutput(
                grads=clipped(zeros, max_norm=output_max_norm),
                squared_grads=clipped(
                    tree_map(
                        lambda x: (
                            torch.zeros_like(x) if isinstance(x, torch.Tensor) else x
                        ),
                        zeros,
                    ),
                    max_norm=output_squared_max_norm,
                ),
            )
        else:
            grads = clipped(zeros, max_norm=output_max_norm)
        if return_aux:
            empty = torch.empty(0)
            # Per-group steps report ``clipping_rate=None`` (rates live in
            # ``group_norms``), empty draws included, so ``sync(aux)`` matches.
            grad_aux = ClippedGradAux(
                loss_values=empty if return_aux else None,
                grad_norms=empty,
                clipped_grad_norms=empty,
                loss_aux=None,
                clipping_rate=None if isinstance(clipping_norm, PerGroup) else 0.0,
                batch_size=0,
            )
            return (grads, grad_aux), state
        if return_stats:
            from opake.api.engine.clipping._clipped_fun import ClippingStats

            if isinstance(clipping_norm, PerGroup):
                group_names = sorted(clipping_norm.values)
                empty_counts = dict.fromkeys(group_names, 0.0)
                empty_rates = dict.fromkeys(group_names, 0.0)
                stats = ClippingStats(
                    num_clipped=empty_counts,
                    clipping_rate=empty_rates,
                    batch_size=0,
                )
            else:
                stats = ClippingStats(num_clipped=0.0, clipping_rate=0.0, batch_size=0)
            return (grads, stats), state
        return grads, state

    # Use PyTorch's grad_and_value (returns (grad, value) or (grad, (value, aux)))
    grad_and_value_fn = grad_and_value(loss_fn, argnums=argnums, has_aux=True)

    def grad_fn(*args, **kwargs):
        # Keep the graph only when an enclosing transform differentiates it.
        values_only = not under_differentiating_transform(when_compiling=True)
        with torch.no_grad() if values_only else contextlib.nullcontext():
            grad, value_and_aux = grad_and_value_fn(*args, **kwargs)
        result = pre_clipping_transform(grad)
        if return_aux:
            # Return dict aux from per-example grad_fn; clipping-related norms are
            # produced by clipped_fun to avoid duplicating norm computation logic.
            # PyTorch vmap cannot handle namedtuples with None values when out_dims != None
            aux_dict = {}
            if return_aux:
                aux_dict["values"] = value_and_aux[0]
            if return_aux and has_aux:
                aux_dict["value_aux"] = value_and_aux[1]
            return result, aux_dict
        return result

    clipped_grad_fn, clip_state = clipped_fun(
        grad_fn,
        has_aux=return_aux,
        batch_argnums=batch_argnums,
        clipping_norm=clipping_norm,
        normalize_by=normalize_by,
        return_aux=return_aux,
        return_stats=return_stats,
        second_moment=second_moment,
        microbatch_size=microbatch_size,
        dtype=dtype,
        compute_dtype=compute_dtype,
        _scale_fn=_scale_fn,
        _builtin_scale=_builtin_scale,
        _chunk_compiler=_chunk_compiler,
    )

    if return_stats:

        def grad_fn_wrapper(*args, state, **kwargs):
            if batch_size_from_args(args, batch_argnums_tuple) == 0:
                return _empty_batch_response(args, state)
            with record_function("opake::clipped_grad"):
                (clipped_grads, stats), returned_state = clipped_grad_fn(
                    *args, state=state, **kwargs
                )
            return (clipped_grads, stats), returned_state

        return grad_fn_wrapper, clip_state

    # clipped_grad_fn is now a callable, clip_state is a FixedClipState
    # Wrap the result to convert ClippedFunAux to ClippedGradAux
    if not return_aux:
        # No aux, return wrapped directly with state-passing signature
        def grad_fn_wrapper(*args, state, **kwargs):
            if batch_size_from_args(args, batch_argnums_tuple) == 0:
                return _empty_batch_response(args, state)

            with record_function("opake::clipped_grad"):
                (result, returned_state) = clipped_grad_fn(*args, state=state, **kwargs)
            return result, returned_state

        return grad_fn_wrapper, clip_state
    else:
        # Need to convert ClippedFunAux to ClippedGradAux
        def grad_fn_wrapper(*args, state, **kwargs):
            if batch_size_from_args(args, batch_argnums_tuple) == 0:
                return _empty_batch_response(args, state)

            with record_function("opake::clipped_grad"):
                (clipped_grads, aux), returned_state = clipped_grad_fn(
                    *args, state=state, **kwargs
                )
            grad_aux = ClippedGradAux(
                loss_values=aux.values,
                grad_norms=aux.norms,
                clipped_grad_norms=aux.clipped_norms,
                loss_aux=aux.value_aux,
                clipping_rate=aux.clipping_rate,
                batch_size=aux.batch_size,
                group_norms=aux.group_norms,
            )
            return (clipped_grads, grad_aux), returned_state

        return grad_fn_wrapper, clip_state

`per_group(params, /, patterns=None, *, fallback=None, allow_unused_patterns=False, **kwargs)` ¶

Construct PerGroup from parameter leaf paths and substring patterns.

Each pattern is a substring matched against the dotted display form of each leaf :data:~opake.pytree.ParamPath (see :func:~opake.pytree.param_path_display). Leaves whose display path contains the substring are assigned to that group. Every leaf must match exactly one pattern (error on 0 or 2+).

params may be any tensor pytree: flat named_parameters dicts (paths are one-segment (name,)), nested dicts, lists, or tuples. The returned :class:~opake.types.PerGroup is compiled against those optree paths so clip / noise / optimizers look up the same identity.

The patterns dict arg merges with **kwargs to support keys containing dots (which cannot be used as keyword arguments)::

per_group(params, **{"layers.0": 0.5, "layers.1": 1.0})
# equivalent to:
per_group(params, patterns={"layers.0": 0.5, "layers.1": 1.0})

The fallback parameter assigns a value to any parameter that does not match an explicit pattern. Without fallback, unmatched parameters raise ValueError.

By default, every explicit pattern must match at least one parameter. Set allow_unused_patterns=True when intentionally reusing a configuration across models with different parameter names.

Parameters:

Name	Type	Description	Default
`params`	`Any`	Parameter pytree (flat or nested). Leaf paths are matched against patterns via their dotted display form.	required
`patterns`	`dict[str, float] \| None`	Optional dict of `{pattern: value}` pairs. Merged with `**kwargs`.	`None`
`fallback`	`float \| None`	Optional value for parameters that don't match any explicit pattern. Unmatched params are assigned to a group named `"fallback"`.	`None`
`allow_unused_patterns`	`bool`	Whether to allow explicit patterns that do not match any parameter. Defaults to `False`.	`False`
`**kwargs`	`float`	Pattern-value pairs where the kwarg name is the substring pattern and the value is the per-group value (e.g. clipping norm).	`{}`

Returns:

Type	Description
`PerGroup`	class:`~opake.types.PerGroup` with pre-resolved
`PerGroup`	path-to-group assignments.

`fun` ¶

DP-FTRL clipping power-user helpers.

:func:clipped_fun / :func:auto_clipped_fun — clip + sum any per-example function output (PyTree).
:func:clip_pytree / :func:auto_scale_pytree — clip / AUTO-S an already-batched pytree.

`auto_clipped_fun(fun, has_aux=False, *, batch_argnums=0, R=1.0, gamma=_DEFAULT_GAMMA, normalize_by=1.0, return_aux=False, return_stats=False, microbatch_size=None, dtype=None, _chunk_compiler=None)` ¶

Transform a function so each per-example output is scaled to clipped norm via AUTO-S.

Mirrors :func:clipped_fun but replaces threshold-based clipping with AUTO-S automatic scaling R \cdot v / (\|v\| + \gamma). The output of the returned function is the sum of the scaled per-example outputs.

Parameters:

Name	Type	Description	Default
`fun`	`Callable`	The per-example function to be transformed.	required
`has_aux`	`bool`	If True, `fun` returns `(value, aux)`; only `value` is scaled and aggregated.	`False`
`batch_argnums`	`int \| tuple[int, ...]`	Which arguments have a batch dimension.	`0`
`R`	`float \| PerGroup`	Sensitivity max_norm for the scaled output. When `PerGroup`, each group is scaled independently.	`1.0`
`gamma`	`float`	Denominator stabilizer :math:`\gamma` (default 0.01).	`_DEFAULT_GAMMA`
`normalize_by`	`float`	Divisor applied to the scaled sum.	`1.0`
`return_aux`	`bool`	If True, the returned callable returns an :class:`AutoClippedFunAux` alongside the summed value.	`False`
`return_stats`	`bool`	If True, return :class:`ClippingStats` alongside the summed value. Cannot be combined with `return_aux`.	`False`
`microbatch_size`	`int \| None`	Process the batch in chunks of this size.	`None`
`dtype`	`Any`	Optional accumulation dtype for the sum.	`None`

Returns:

Type	Description
`Callable`	`(auto_fn, state)`. `auto_fn` has the signature
`AutoClipState`	`(args, state, *kwargs) -> (value, state)` or
`tuple[Callable, AutoClipState]`	`(args, state, *kwargs) -> ((value, aux), state)` when
`tuple[Callable, AutoClipState]`	`return_aux=True`.

Raises:

Type	Description
`ConfigurationError`	If the batch is empty.

Formal guarantee

Under add/remove or zero-out DP, the L2 sensitivity of the first output with respect to the batch arguments is the returned ClippedPytree.max_norm metadata. The bound is constant and data-independent, so per-step Gaussian calibration is correct. mf_gaussian_noise additionally requires that max_norm not drift across steps; AUTO-S satisfies that latch because R and normalize_by are fixed.

Source code in packages/opake-engine/src/opake/api/engine/clipping/_auto.py

def auto_clipped_fun(
    fun: Callable,
    has_aux: bool = False,
    *,
    batch_argnums: int | tuple[int, ...] = 0,
    R: float | PerGroup = 1.0,
    gamma: float = _DEFAULT_GAMMA,
    normalize_by: float = 1.0,
    return_aux: bool = False,
    return_stats: bool = False,
    microbatch_size: int | None = None,
    dtype: Any = None,
    _chunk_compiler: Callable | None = None,
) -> tuple[Callable, AutoClipState]:
    r"""Transform a function so each per-example output is scaled to clipped norm via AUTO-S.

    Mirrors :func:`clipped_fun` but replaces threshold-based clipping with
    AUTO-S automatic scaling ``R \cdot v / (\|v\| + \gamma)``.  The output of
    the returned function is the sum of the scaled per-example outputs.

    Args:
        fun: The per-example function to be transformed.
        has_aux: If True, ``fun`` returns ``(value, aux)``; only ``value``
            is scaled and aggregated.
        batch_argnums: Which arguments have a batch dimension.
        R: Sensitivity max_norm for the scaled output.  When ``PerGroup``,
            each group is scaled independently.
        gamma: Denominator stabilizer :math:`\gamma` (default 0.01).
        normalize_by: Divisor applied to the scaled sum.
        return_aux: If True, the returned callable returns an
            :class:`AutoClippedFunAux` alongside the summed value.
        return_stats: If True, return :class:`ClippingStats` alongside the
            summed value. Cannot be combined with ``return_aux``.
        microbatch_size: Process the batch in chunks of this size.
        dtype: Optional accumulation dtype for the sum.

    Returns:
        ``(auto_fn, state)``.  ``auto_fn`` has the signature
        ``(*args, state, **kwargs) -> (value, state)`` or
        ``(*args, state, **kwargs) -> ((value, aux), state)`` when
        ``return_aux=True``.

    Raises:
        ConfigurationError: If the batch is empty.

    Formal guarantee:
        Under add/remove or zero-out DP, the L2 sensitivity of the first
        output with respect to the batch arguments is the returned
        ``ClippedPytree.max_norm`` metadata.  The bound is constant and
        data-independent, so per-step Gaussian calibration is correct.
        ``mf_gaussian_noise`` additionally requires that ``max_norm`` not drift across
        steps; AUTO-S satisfies that latch because ``R`` and ``normalize_by``
        are fixed.
    """
    _validate_auto_params(R, gamma)

    scale_fn = _make_auto_scale_fn(R, gamma)
    builtin_scale = _BuiltinScale(kind="auto_s", gamma=gamma)
    inner_fn, _ = clipped_fun(
        fun,
        has_aux=has_aux,
        batch_argnums=batch_argnums,
        clipping_norm=R,
        normalize_by=normalize_by,
        return_aux=return_aux,
        return_stats=return_stats,
        microbatch_size=microbatch_size,
        dtype=dtype,
        _scale_fn=scale_fn,
        _builtin_scale=builtin_scale,
        _chunk_compiler=_chunk_compiler,
    )

    state = AutoClipState()

    if not return_aux:

        def auto_fn(*args, state, **kwargs):
            result, _ = inner_fn(*args, state=None, **kwargs)
            return result, state

        return auto_fn, state

    def auto_fn(*args, state, **kwargs):
        (result, fun_aux), _ = inner_fn(*args, state=None, **kwargs)
        aux = AutoClippedFunAux(
            values=fun_aux.values,
            norms=fun_aux.norms,
            clipped_norms=fun_aux.clipped_norms,
            value_aux=fun_aux.value_aux,
            clipping_rate=fun_aux.clipping_rate,
            batch_size=fun_aux.batch_size,
            group_norms=fun_aux.group_norms,
        )
        return (result, aux), state

    return auto_fn, state

`auto_scale_pytree(pytree, R=1.0, gamma=0.01, *, compute_dtype=None)` ¶

AUTO-S automatic scaling of a PyTree (Bu et al., NeurIPS 2023).

Scales the PyTree by R / (\|pytree\| + gamma) so the output L2 norm is clipped by R for any input. Unlike :func:clip_pytree, there is no threshold to tune — every example contributes an approximately unit-length update, with the effective step size absorbed into the learning rate.

Parameters:

Name	Type	Description	Default
`pytree`	`Any`	Tensor pytree to scale (flat or nested).	required
`R`	`float \| PerGroup`	Output sensitivity bound (non-negative). When `PerGroup`, each group is scaled independently to its own bound; leaf paths must match `R.groups` exactly.	`1.0`
`gamma`	`float`	Small positive denominator stabilizer :math:`\gamma` (default 0.01). Must be strictly positive; at `gamma=0` this reduces to AUTO-V (undefined at zero gradient).	`0.01`
`compute_dtype`	`dtype \| None`	Internal dtype for the L2-norm reduction and leaf scaling. `None` (default) auto-promotes bf16/fp16 inputs to float32.	`None`

Returns:

Type	Description
`Any`	Tuple of (scaled_pytree, aux) where `aux.norm` is the original L2
`ClipPytreeAux`	norm and `aux.group_norms` carries per-group norms in per-group
`tuple[Any, ClipPytreeAux]`	mode.

Formal guarantee

For every input, the output has L2 norm at most R (scalar case) or :math:\sqrt{\sum_k R_k^2} (per-group case). The bound holds on the values as stored: each leaf is scaled in compute_dtype (never below its storage precision) and cast once to its output dtype. The scale is shrunk by the final cast's rounding margin so the result cannot exceed R. Unlike :func:clip_pytree there is no min(1, ...), so the shrink applies to every input — roughly 0.4% of the output magnitude at bfloat16, 0.05% at float16, and a few ULPs at float32.

NaN/Inf values are sanitized to zero before scaling, matching the behavior of :func:clip_pytree.

Source code in packages/opake-engine/src/opake/api/engine/clipping/_pytree.py

def auto_scale_pytree(
    pytree: Any,
    R: float | PerGroup = 1.0,
    gamma: float = 0.01,
    *,
    compute_dtype: torch.dtype | None = None,
) -> tuple[Any, ClipPytreeAux]:
    r"""AUTO-S automatic scaling of a PyTree (Bu et al., NeurIPS 2023).

    Scales the PyTree by ``R / (\|pytree\| + gamma)`` so the output L2
    norm is clipped by ``R`` for any input. Unlike :func:`clip_pytree`, there
    is no threshold to tune — every example contributes an approximately
    unit-length update, with the effective step size absorbed into the
    learning rate.

    Args:
        pytree: Tensor pytree to scale (flat or nested).
        R: Output sensitivity bound (non-negative). When ``PerGroup``, each
            group is scaled independently to its own bound; leaf paths must
            match ``R.groups`` exactly.
        gamma: Small positive denominator stabilizer :math:`\gamma` (default
            0.01). Must be strictly positive; at ``gamma=0`` this reduces to
            AUTO-V (undefined at zero gradient).
        compute_dtype: Internal dtype for the L2-norm reduction and leaf
            scaling. ``None`` (default) auto-promotes bf16/fp16 inputs to
            float32.

    Returns:
        Tuple of (scaled_pytree, aux) where ``aux.norm`` is the original L2
        norm and ``aux.group_norms`` carries per-group norms in per-group
        mode.

    Formal guarantee:
        For every input, the output has L2 norm at most ``R`` (scalar case)
        or :math:`\sqrt{\sum_k R_k^2}` (per-group case).  The bound holds on
        the values as *stored*: each leaf is scaled in ``compute_dtype`` (never
        below its storage precision) and cast once to its output dtype.  The
        scale is shrunk by the final cast's rounding margin so the result cannot
        exceed ``R``. Unlike :func:`clip_pytree` there is no ``min(1, ...)``, so
        the shrink applies to every input — roughly 0.4% of the output magnitude
        at bfloat16, 0.05% at float16, and a few ULPs at float32.

    NaN/Inf values are sanitized to zero before scaling, matching the
    behavior of :func:`clip_pytree`.
    """
    if gamma <= 0:
        raise ConfigurationError(*(f"gamma must be positive, got {gamma}",))

    pytree = tree_map(
        lambda t: (
            torch.nan_to_num(t, nan=0.0, posinf=0.0, neginf=0.0)
            if isinstance(t, torch.Tensor)
            else t
        ),
        pytree,
    )

    if isinstance(R, PerGroup):
        return _auto_scale_per_group(pytree, R, gamma, compute_dtype)

    leaves = [leaf for leaf in tree_leaves(pytree) if isinstance(leaf, torch.Tensor)]
    acc_dtype = _resolve_compute_dtype_for_reduction(leaves, compute_dtype)
    sq_dtype = _sq_accum_dtype(leaves)
    roundoff = _norm_roundoff(
        acc_dtype, sq_dtype, len(leaves), _reduction_terms(leaves)
    )

    norm = torch.sqrt(_sq_norm(leaves, acc_dtype, sq_dtype))
    orig_norm = norm.to(acc_dtype)

    R_tensor = torch.clamp(
        torch.as_tensor(R, dtype=norm.dtype, device=norm.device), min=0.0
    )
    gamma_tensor = torch.tensor(gamma, dtype=norm.dtype, device=norm.device)
    ratio = R_tensor / (norm + gamma_tensor)

    def scale_leaf(t):
        if not isinstance(t, torch.Tensor):
            return t
        return _scale_tensor(t, ratio, acc_dtype, roundoff, clamp_to_one=False)

    scaled = tree_map(
        lambda t: scale_leaf(t) if isinstance(t, torch.Tensor) else t, pytree
    )

    return scaled, ClipPytreeAux(norm=orig_norm, group_norms=None)

`clip_pytree(pytree, clipping_norm, return_zero=False, *, compute_dtype=None)` ¶

Clip a PyTree of tensors to a maximum L2 norm.

NaN and Inf values in the input are replaced with zeros before clipping. This is vmap-compatible and DP-safe: norm(output) <= clipping_norm holds on the values as stored, not just in exact arithmetic. The norm reduction, the multiplication run in compute_dtype (never below the leaf's storage precision), and the result is cast once to the leaf dtype. That cast can round up, so each leaf's scale is shrunk by its final-rounding margin. The shrink is applied before the min(1, ...), so an input well inside the threshold is returned unchanged; one whose norm sits within a few ULPs of clipping_norm is scaled by that much. For clipped inputs the cost is roughly 0.4% of the output magnitude at bfloat16, 0.05% at float16, and a few ULPs at float32.

Parameters:

Name	Type	Description	Default
`pytree`	`Any`	Tensor pytree to clip (flat or nested).	required
`clipping_norm`	`float \| PerGroup`	Maximum L2 norm (non-negative, or inf for no clipping). When `PerGroup`, each group is clipped independently; leaf paths must match `clipping_norm.groups` exactly.	required
`return_zero`	`bool`	If True, the output PyTree is guaranteed to be zero no matter what the inputs are. Does not influence the formal guarantees but useful for privacy amplification via padding (see https://arxiv.org/pdf/2411.04205).	`False`
`compute_dtype`	`dtype \| None`	Internal dtype for the L2-norm reduction and leaf scaling. `None` (default) auto-promotes bf16/fp16 inputs to float32. This keeps the sensitivity bound numerically honest under low-precision compute — small per-element contributions don't get rounded away during the sum-of-squares or multiplication.	`None`

Returns:

Type	Description
`tuple[Any, ClipPytreeAux]`	Tuple of (clipped_pytree, aux) where aux contains: - norm: The L2 norm of the original (unclipped) pytree

Edge cases

clipping_norm=0: Returns zeros
clipping_norm=inf: No clipping (passthrough)
pytree_norm=0: Returns unchanged
NaN/Inf values: Replaced with zeros before clipping
return_zero=True: Returns zeros regardless of other parameters

Source code in packages/opake-engine/src/opake/api/engine/clipping/_pytree.py

def clip_pytree(
    pytree: Any,
    clipping_norm: float | PerGroup,
    return_zero: bool = False,
    *,
    compute_dtype: torch.dtype | None = None,
) -> tuple[Any, ClipPytreeAux]:
    """Clip a PyTree of tensors to a maximum L2 norm.

    NaN and Inf values in the input are replaced with zeros before clipping.
    This is vmap-compatible and DP-safe: ``norm(output) <= clipping_norm`` holds
    on the values as *stored*, not just in exact arithmetic.  The norm reduction,
    the multiplication run in ``compute_dtype`` (never below the leaf's storage
    precision), and the result is cast once to the leaf dtype. That cast can
    round up, so each leaf's scale is shrunk by its final-rounding margin. The
    shrink is applied before the ``min(1, ...)``, so an input well inside the
    threshold is returned unchanged; one whose norm sits within a few ULPs of
    ``clipping_norm`` is scaled by that much. For clipped inputs the cost is
    roughly 0.4% of the output magnitude at bfloat16, 0.05% at float16, and a
    few ULPs at float32.

    Args:
        pytree: Tensor pytree to clip (flat or nested).
        clipping_norm: Maximum L2 norm (non-negative, or inf for no clipping).
            When ``PerGroup``, each group is clipped independently; leaf paths
            must match ``clipping_norm.groups`` exactly.
        return_zero: If True, the output PyTree is guaranteed to be zero no matter
            what the inputs are. Does not influence the formal guarantees but useful
            for privacy amplification via padding (see https://arxiv.org/pdf/2411.04205).
        compute_dtype: Internal dtype for the L2-norm reduction and leaf
            scaling. ``None`` (default) auto-promotes bf16/fp16 inputs to
            float32. This keeps the sensitivity bound numerically honest under
            low-precision compute — small per-element contributions don't get
            rounded away during the sum-of-squares or multiplication.

    Returns:
        Tuple of (clipped_pytree, aux) where aux contains:
            - norm: The L2 norm of the original (unclipped) pytree

    Edge cases:
        - clipping_norm=0: Returns zeros
        - clipping_norm=inf: No clipping (passthrough)
        - pytree_norm=0: Returns unchanged
        - NaN/Inf values: Replaced with zeros before clipping
        - return_zero=True: Returns zeros regardless of other parameters
    """
    pytree = tree_map(
        lambda t: (
            torch.nan_to_num(t, nan=0.0, posinf=0.0, neginf=0.0)
            if isinstance(t, torch.Tensor)
            else t
        ),
        pytree,
    )

    if isinstance(clipping_norm, PerGroup):
        return _clip_pytree_per_group(pytree, clipping_norm, return_zero, compute_dtype)

    leaves = [leaf for leaf in tree_leaves(pytree) if isinstance(leaf, torch.Tensor)]
    acc_dtype = _resolve_compute_dtype_for_reduction(leaves, compute_dtype)
    sq_dtype = _sq_accum_dtype(leaves)
    roundoff = _norm_roundoff(
        acc_dtype, sq_dtype, len(leaves), _reduction_terms(leaves)
    )

    sq_norm = _sq_norm(leaves, acc_dtype, sq_dtype)
    norm = torch.sqrt(sq_norm)
    orig_norm = norm.to(acc_dtype)

    clipping_norm_tensor = torch.as_tensor(
        clipping_norm, dtype=norm.dtype, device=norm.device
    )
    clipping_norm_tensor = torch.clamp(clipping_norm_tensor, min=0.0)
    ratio = clipping_norm_tensor / norm

    def scale_leaf(t):
        if not isinstance(t, torch.Tensor):
            return t
        return _scale_tensor(t, ratio, acc_dtype, roundoff, clamp_to_one=True)

    clipped = tree_map(
        lambda t: scale_leaf(t) if isinstance(t, torch.Tensor) else t, pytree
    )

    if return_zero:
        clipped = tree_map(
            lambda t: torch.zeros_like(t) if isinstance(t, torch.Tensor) else t, clipped
        )

    return clipped, ClipPytreeAux(norm=orig_norm, group_norms=None)

`clipped_fun(fun, has_aux=False, *, batch_argnums=0, clipping_norm=1.0, normalize_by=1.0, return_aux=False, return_stats=False, second_moment=False, microbatch_size=None, dtype=None, compute_dtype=None, _scale_fn=None, _builtin_scale=None, _chunk_compiler=None)` ¶

Transform a function to clip its output and sum across a batch.

This is the primary API for per-example clipping in DP-SGD. It wraps a function to clip each per-example output to a maximum L2 norm, then sums the clipped outputs.

The returned pytree is wrapped as :class:ClippedPytree (single-stream) or :class:SecondMomentClippingOutput (paired-stream when second_moment=True), carrying the post-normalization max_norm for downstream noise calibration. The bound is part of the contract: consumers (gaussian_noise, mf_gaussian_noise) read it directly without the caller threading a separate sensitivity argument. Unwrap to a raw pytree via .pytree if you need the summed values without metadata.

Example Usage

from opake.api.engine.clipping._clipped_fun import clipped_fun data = torch.tensor([0.0, 1.0, 2.0, 3.0, 4.0, 5.0]) clipped_mean, clip_state = clipped_fun(torch.mean, clipping_norm=1.0) result, clip_state = clipped_mean(data, state=clip_state) result.pytree tensor(5.) result.max_norm 1.0

Formal Guarantees

For the first function output: The L2 sensitivity of the returned function with respect to the batch arguments (specified by batch_argnums) under add/remove or zero-out differential privacy definitions is guaranteed to be clipping_norm. Under replace-one DP, the sensitivity is doubled (2 * clipping_norm). Extra auxiliary outputs (aux, norms) are per-example. This function guarantees that per-example outputs only depend on the data for the same example. This allows maximum flexibility for the caller to aggregate these as desired (possibly with a DP mean, median, quantile, or histogram mechanism).

Parameters:

Name	Type	Description	Default
`fun`	`Callable[..., Any]`	The function to be clipped.	required
`has_aux`	`bool`	If True, `fun` is expected to return a tuple `(value, loss_aux)`. Only the value will be clipped + aggregated, `loss_aux` will be returned on a per-example basis. Exercise caution when using this as the sensitivity guarantees of the returned Callable are only provided w.r.t. `value`.	`False`
`batch_argnums`	`int \| tuple[int, ...]`	Specifies which argument(s) of `fun` contain the batch dimension. All arguments specified here must have the same size along the 0th axis.	`0`
`clipping_norm`	`float \| PerGroup`	The maximum L2 norm allowed.	`1.0`
`normalize_by`	`float`	Divide the clipped output by this value before returning.	`1.0`
`return_aux`	`bool`	If True, the returned Callable will return a per-example aux dataclass containing the original per-example values, per-example norms before clipping, and any auxiliary data returned by `fun`.	`False`
`return_stats`	`bool`	If True, return :class:`ClippingStats` with aggregate clipping counts, rates, and batch size. Cannot be combined with `return_aux`.	`False`
`second_moment`	`bool`	If True, also accumulate the element-wise sum of per-example squared clipped values, i.e. `Σᵢ gᵢ²`. The squaring happens inside the per-example loop so the second-stream sensitivity is the per-record squared bound `C²` (averaged: `C² / normalize_by`). The wrapped output becomes :class:`SecondMomentClippingOutput` with both streams. Per-group `clipping_norm` is supported and yields `SecondMomentClippingOutput` with per-group `max_norm` on both streams.	`False`
`microbatch_size`	`int \| None`	If set, the batch is split up into microbatches of this size for memory-efficient processing. Processes each microbatch separately and accumulates results without materializing the full batch of gradients. Set this to reduce peak memory usage at the cost of slightly slower computation. The running sum is held in the wider of `compute_dtype` and the output dtype, so a bf16/fp16 run keeps one float32 copy of the summed output. Pass `compute_dtype=torch.bfloat16` to give that memory back, at the cost of accumulation precision when the output dtype is also low precision.	`None`
`dtype`	`dtype \| None`	Optional dtype for the clipped+aggregated pytree. If None, the dtype will be the same as the dtypes of the function output.	`None`
`compute_dtype`	`dtype \| None`	Internal reduction and accumulation dtype. `None` (default) auto-promotes bf16/fp16 to float32 for numerical stability; an explicit dtype selects the reduction precision. Leaf scaling and microbatch accumulation use the wider of this dtype and the leaf or output storage dtype, so explicit lower precision never narrows a wider value. Independent of `dtype` (which controls the output dtype). Applies across microbatches too, so microbatched and non-microbatched runs agree at their resolved accumulation precision.	`None`

Returns: A tuple (clip_fn, FixedClipState) where clip_fn(*args, state=...) clips the output of fun and sums across the batch. The exact return shape depends on second_moment and return_aux:

| ``second_moment`` | ``return_aux`` | ``clip_fn`` returns                                |
| :---------------- | :------------- | :------------------------------------------------- |
| False             | False          | ``(ClippedPytree, state)``                         |
| False             | True           | ``((ClippedPytree, ClippedFunAux), state)``        |
| True              | False          | ``(SecondMomentClippingOutput, state)``            |
| True              | True           | ``((SecondMomentClippingOutput, ClippedFunAux), state)`` |

With ``return_stats=True`` and ``return_aux=False``, the clipped value
is paired with :class:`ClippingStats`.

Raises:

Type	Description
`ConfigurationError`	If the batch is empty or a chunk compiler is used without a finite microbatch size.

Source code in packages/opake-engine/src/opake/api/engine/clipping/_clipped_fun.py

def clipped_fun(
    fun: Callable[..., Any],
    has_aux: bool = False,
    *,
    batch_argnums: int | tuple[int, ...] = 0,
    clipping_norm: float | PerGroup = 1.0,
    normalize_by: float = 1.0,
    return_aux: bool = False,
    return_stats: bool = False,
    second_moment: bool = False,
    microbatch_size: int | None = None,
    dtype: torch.dtype | None = None,
    compute_dtype: torch.dtype | None = None,
    _scale_fn: Callable | None = None,
    _builtin_scale: _BuiltinScale | None = None,
    _chunk_compiler: Callable | None = None,
) -> tuple[Callable, FixedClipState]:
    """Transform a function to clip its output and sum across a batch.

    This is the primary API for per-example clipping in DP-SGD. It wraps a function
    to clip each per-example output to a maximum L2 norm, then sums the clipped outputs.

    The returned pytree is wrapped as :class:`ClippedPytree` (single-stream)
    or :class:`SecondMomentClippingOutput` (paired-stream when
    ``second_moment=True``), carrying the post-normalization
    ``max_norm`` for downstream noise calibration.  The bound is part of
    the contract: consumers (``gaussian_noise``, ``mf_gaussian_noise``) read it
    directly without the caller threading a separate ``sensitivity``
    argument.  Unwrap to a raw pytree via ``.pytree`` if you need the
    summed values without metadata.

    Example Usage:
        >>> from opake.api.engine.clipping._clipped_fun import clipped_fun
        >>> data = torch.tensor([0.0, 1.0, 2.0, 3.0, 4.0, 5.0])
        >>> clipped_mean, clip_state = clipped_fun(torch.mean, clipping_norm=1.0)
        >>> result, clip_state = clipped_mean(data, state=clip_state)
        >>> result.pytree
        tensor(5.)
        >>> result.max_norm
        1.0

    Formal Guarantees:
        For the first function output:
          The L2 sensitivity of the returned function with respect to the batch
          arguments (specified by `batch_argnums`) under add/remove or zero-out
          differential privacy definitions is guaranteed to be `clipping_norm`.
          Under replace-one DP, the sensitivity is doubled (2 * `clipping_norm`).
        Extra auxiliary outputs (aux, norms) are per-example. This function
          guarantees that per-example outputs only depend on the data for the same
          example. This allows maximum flexibility for the caller to aggregate
          these as desired (possibly with a DP mean, median, quantile, or histogram
          mechanism).

    Args:
        fun: The function to be clipped.
        has_aux: If True, `fun` is expected to return a tuple `(value, loss_aux)`. Only
            the value will be clipped + aggregated, `loss_aux` will be returned on a
            per-example basis. Exercise caution when using this as the sensitivity
            guarantees of the returned Callable are only provided w.r.t. `value`.
        batch_argnums: Specifies which argument(s) of `fun` contain the batch
            dimension. All arguments specified here must have the same size along the
            0th axis.
        clipping_norm: The maximum L2 norm allowed.
        normalize_by: Divide the clipped output by this value before returning.
        return_aux: If True, the returned Callable will return a per-example aux
            dataclass containing the original per-example values, per-example norms
            before clipping, and any auxiliary data returned by `fun`.
        return_stats: If True, return :class:`ClippingStats` with aggregate
            clipping counts, rates, and batch size. Cannot be combined with
            ``return_aux``.
        second_moment: If True, also accumulate the element-wise sum of
            per-example squared clipped values, i.e. ``Σᵢ gᵢ²``.  The
            squaring happens inside the per-example loop so the
            second-stream sensitivity is the per-record squared bound
            ``C²`` (averaged: ``C² / normalize_by``).  The wrapped
            output becomes :class:`SecondMomentClippingOutput` with both
            streams.  Per-group ``clipping_norm`` is supported and yields
            ``SecondMomentClippingOutput`` with per-group ``max_norm``
            on both streams.
        microbatch_size: If set, the batch is split up into microbatches of this
            size for memory-efficient processing. Processes each microbatch separately
            and accumulates results without materializing the full batch of gradients.
            Set this to reduce peak memory usage at the cost of slightly slower computation.
            The running sum is held in the wider of ``compute_dtype`` and the
            output dtype, so a bf16/fp16 run keeps one float32 copy of the summed
            output. Pass ``compute_dtype=torch.bfloat16`` to give that memory back,
            at the cost of accumulation precision when the output dtype is also low
            precision.
        dtype: Optional dtype for the clipped+aggregated pytree. If None, the dtype
            will be the same as the dtypes of the function output.
        compute_dtype: Internal reduction and accumulation dtype. ``None``
            (default) auto-promotes bf16/fp16 to float32 for numerical stability;
            an explicit dtype selects the reduction precision. Leaf scaling and
            microbatch accumulation use the wider of this dtype and the leaf or
            output storage dtype, so explicit lower precision never narrows a wider
            value. Independent of ``dtype`` (which controls the *output* dtype).
            Applies across microbatches too, so microbatched and non-microbatched
            runs agree at their resolved accumulation precision.
    Returns:
        A tuple ``(clip_fn, FixedClipState)`` where ``clip_fn(*args, state=...)``
        clips the output of ``fun`` and sums across the batch.  The exact
        return shape depends on ``second_moment`` and ``return_aux``:

        | ``second_moment`` | ``return_aux`` | ``clip_fn`` returns                                |
        | :---------------- | :------------- | :------------------------------------------------- |
        | False             | False          | ``(ClippedPytree, state)``                         |
        | False             | True           | ``((ClippedPytree, ClippedFunAux), state)``        |
        | True              | False          | ``(SecondMomentClippingOutput, state)``            |
        | True              | True           | ``((SecondMomentClippingOutput, ClippedFunAux), state)`` |

        With ``return_stats=True`` and ``return_aux=False``, the clipped value
        is paired with :class:`ClippingStats`.

    Raises:
        ConfigurationError: If the batch is empty or a chunk compiler is used
            without a finite microbatch size.
    """
    batch_argnums = _prepare_clipped_fun(
        batch_argnums, clipping_norm, return_aux, return_stats
    )
    stream_impl = _resolve_stream_impl(
        clipping_norm=clipping_norm,
        second_moment=second_moment,
        compute_dtype=compute_dtype,
        scale_fn=_scale_fn,
        chunk_compiler=_chunk_compiler,
    )

    # Wrap function to handle has_aux - use empty tuple () not None!
    if not has_aux:

        def fun_with_aux(*args, **kwargs):
            return (fun(*args, **kwargs), ())

    else:
        fun_with_aux = fun

    clip_state = FixedClipState()

    def _per_example_fn(kernel_clipping_norm, *args_single, **call_kwargs):
        value, aux = fun_with_aux(*args_single, **call_kwargs)
        if _scale_fn is None:
            clipped_value, norm = clip_pytree(
                value,
                clipping_norm=kernel_clipping_norm,
                compute_dtype=compute_dtype,
            )
        else:
            clipped_value, norm = _scale_fn(value)
        squared_value = (
            tree_map(
                lambda x: x.square() if isinstance(x, torch.Tensor) else x,
                clipped_value,
            )
            if second_moment
            else None
        )
        if return_aux or return_stats:
            diagnostics = {"norms": norm.norm.detach()}
            if return_aux:
                diagnostics["clipped_norms"] = global_norm(
                    clipped_value, compute_dtype=compute_dtype
                ).detach()
            if norm.group_norms is not None:
                diagnostics["group_norms"] = {
                    key: group_norm.detach()
                    for key, group_norm in norm.group_norms.items()
                }

            if return_aux:
                _attach_value_aux(diagnostics, value, aux, has_aux)

            if second_moment:
                return clipped_value, squared_value, diagnostics
            return clipped_value, diagnostics
        if second_moment:
            return clipped_value, squared_value
        return clipped_value

    if _chunk_compiler is not None and microbatch_size is None:
        raise ConfigurationError(
            *("_chunk_compiler requires a finite microbatch_size",)
        )

    def _streaming_kernel(in_dims, kernel_clipping_norm, args, kwargs, reduce_leaf):
        values, value_aux = _vmap_values(
            fun_with_aux, in_dims, args, kwargs, return_aux
        )
        reduced, markers, squared_reduced, squared_markers, diagnostics = (
            _stream_clip_and_sum if stream_impl is None else stream_impl
        )(
            values,
            kernel_clipping_norm,
            scale=_FIXED_SCALE if _scale_fn is None else _builtin_scale,
            batch_size=batch_size_from_args(args, batch_argnums),
            reduce_leaf=reduce_leaf,
            compute_dtype=compute_dtype,
            second_moment=second_moment,
            return_aux=return_aux,
            return_stats=return_stats,
        )
        if return_aux:
            _attach_value_aux(diagnostics, values, value_aux, has_aux)
        return reduced, markers, squared_reduced, squared_markers, diagnostics

    def _make_chunk_kernel(in_dims, stream):
        # Fixed and AUTO-S scaling are construction-time constants.  A
        # stateful scaling parameter must cross this private boundary as a
        # tensor input rather than be captured as changing Python state.
        def _chunk_kernel(kernel_clipping_norm, *chunk_args, **call_kwargs):
            """Run vmap, clipping, diagnostics, and reduction for one chunk."""
            if stream:
                return _streaming_kernel(
                    in_dims,
                    kernel_clipping_norm,
                    chunk_args,
                    call_kwargs,
                    lambda x: torch.sum(
                        x, dim=0, dtype=_accumulation_dtype(x, dtype, compute_dtype)
                    ),
                )
            n_outputs = 1 + int(bool(second_moment)) + int(return_aux or return_stats)
            out_dims = 0 if n_outputs == 1 else (0,) * n_outputs
            vmapped = _vmap(
                _per_example_fn,
                in_dims=(None, *in_dims),
                out_dims=out_dims,
                randomness="same",
            )
            outputs = vmapped(kernel_clipping_norm, *chunk_args, **call_kwargs)
            if n_outputs == 1:
                clipped_values = outputs
                squared_values = None
                diagnostics = ()
            else:
                output_index = 0
                clipped_values = outputs[output_index]
                output_index += 1
                squared_values = outputs[output_index] if second_moment else None
                if second_moment:
                    output_index += 1
                diagnostics = (
                    outputs[output_index] if (return_aux or return_stats) else ()
                )

            reduced = tree_map(
                lambda x: torch.sum(
                    x,
                    dim=0,
                    dtype=_accumulation_dtype(x, dtype, compute_dtype),
                ),
                clipped_values,
            )
            dtype_markers = tree_map(lambda x: x.new_zeros(()), clipped_values)
            if second_moment:
                squared_reduced = tree_map(
                    lambda x: torch.sum(
                        x,
                        dim=0,
                        dtype=_accumulation_dtype(x, dtype, compute_dtype),
                    ),
                    squared_values,
                )
                squared_dtype_markers = tree_map(
                    lambda x: x.new_zeros(()), squared_values
                )
            else:
                squared_reduced = ()
                squared_dtype_markers = ()
            return (
                reduced,
                dtype_markers,
                squared_reduced,
                squared_dtype_markers,
                diagnostics,
            )

        return _chunk_kernel

    # Cache by positional batching structure so callers may use any valid
    # positional/default-argument arrangement without reusing incompatible
    # vmap ``in_dims``.
    chunk_kernels: dict[tuple[tuple[int | None, ...], bool], Callable] = {}

    def clipped_fn(runtime_clipping_norm, *args, **kwargs):
        current_clipping_norm, output_max_norm, output_squared_max_norm = (
            _resolve_runtime_clipping_norm(
                clipping_norm,
                runtime_clipping_norm,
                normalize_by,
                second_moment,
            )
        )
        kernel_clipping_norm = _prepare_kernel_clipping_norm(
            current_clipping_norm, args, batch_argnums
        )
        in_dims = tuple(0 if i in batch_argnums else None for i in range(len(args)))
        stream = not under_differentiating_transform(
            when_compiling=True
        ) and _streaming_supported(current_clipping_norm, _scale_fn, _builtin_scale)

        # Choose execution path based on microbatch_size
        stats = None
        if microbatch_size is None and stream:
            result, _, squared_result, _, aux = _streaming_kernel(
                in_dims,
                kernel_clipping_norm,
                args,
                kwargs,
                lambda x: _sum_clipped_tensor(
                    x, dim=0, output_dtype=dtype, compute_dtype=compute_dtype
                ),
            )
        elif microbatch_size is None:
            # Fast path: vmap entire batch at once.  Output shape depends
            # on the (second_moment, return_aux) flags — see the per_example_fn
            # branches above.  When n_outputs == 1, vmap returns the single
            # pytree (which may itself be a tuple of tensors for tuple
            # params); when n_outputs > 1 the per_example_fn returns a
            # tuple of n_outputs pytrees.
            n_outputs = 1 + int(bool(second_moment)) + int(return_aux or return_stats)
            out_dims = 0 if n_outputs == 1 else (0,) * n_outputs
            vmapped = _vmap(
                _per_example_fn,
                in_dims=(None, *in_dims),
                out_dims=out_dims,
                randomness="same",
            )
            outputs = vmapped(kernel_clipping_norm, *args, **kwargs)
            if n_outputs == 1:
                clipped_values = outputs
                squared_values = None
                aux = ()
            else:
                idx = 0
                clipped_values = outputs[idx]
                idx += 1
                squared_values = outputs[idx] if second_moment else None
                if second_moment:
                    idx += 1
                aux = outputs[idx] if (return_aux or return_stats) else ()

            # Sum clipped values across batch dimension
            result = tree_map(
                lambda x: _sum_clipped_tensor(
                    x, dim=0, output_dtype=dtype, compute_dtype=compute_dtype
                ),
                clipped_values,
            )
            squared_result = (
                tree_map(
                    lambda x: _sum_clipped_tensor(
                        x, dim=0, output_dtype=dtype, compute_dtype=compute_dtype
                    ),
                    squared_values,
                )
                if second_moment
                else None
            )
        else:
            kernel_key = (in_dims, stream)
            chunk_kernel = chunk_kernels.get(kernel_key)
            if chunk_kernel is None:
                eager_chunk_kernel = _make_chunk_kernel(in_dims, stream)
                if _chunk_compiler is None:
                    chunk_kernel = eager_chunk_kernel
                else:
                    # Compilation policy belongs to the owner; the kernel
                    # itself remains tensor-only. Static chunk variants bound
                    # recompilation by ``microbatch_size`` rather than by every
                    # realized Poisson batch size.
                    chunk_kernel = _chunk_compiler(eager_chunk_kernel)
                chunk_kernels[kernel_key] = chunk_kernel
            result, squared_result, aux, stats = _microbatch_accumulate_reduced(
                chunk_fn=chunk_kernel,
                kernel_clipping_norm=kernel_clipping_norm,
                args=args,
                kwargs=kwargs,
                batch_argnums=batch_argnums,
                microbatch_size=microbatch_size,
                return_aux=return_aux,
                return_stats=return_stats,
                dtype=dtype,
                clipping_norm=current_clipping_norm,
                second_moment=second_moment,
            )

        # Normalize
        if normalize_by != 1.0:
            result = tree_map(lambda x: x / normalize_by, result)
            if second_moment:
                squared_result = tree_map(lambda x: x / normalize_by, squared_result)

        if second_moment:
            output = SecondMomentClippingOutput(
                grads=ClippedPytree(pytree=result, max_norm=output_max_norm),
                squared_grads=ClippedPytree(
                    pytree=squared_result, max_norm=output_squared_max_norm
                ),
            )
        else:
            output = clipped(result, max_norm=output_max_norm)

        if not return_aux and not return_stats:
            return output

        if return_aux:
            aux_dict = aux if isinstance(aux, dict) else {}
            norms = aux_dict.get("norms")
            group_norms_dict = aux_dict.get("group_norms")
            stats = _compute_clipping_stats(
                norms,
                clipping_norm=current_clipping_norm,
                group_norms_dict=group_norms_dict,
            )

            aux = ClippedFunAux(
                values=aux_dict.get("values"),
                norms=norms,
                clipped_norms=aux_dict.get("clipped_norms"),
                value_aux=aux_dict.get("value_aux"),
                clipping_rate=(
                    stats.clipping_rate
                    if isinstance(stats.clipping_rate, float)
                    or stats.clipping_rate is None
                    else None
                ),
                batch_size=stats.batch_size,
                group_norms=aux_dict.get("group_norms"),
            )
            return output, aux

        if microbatch_size is not None:
            assert stats is not None
            return output, stats

        aux_dict = aux if isinstance(aux, dict) else {}
        return output, _compute_clipping_stats(
            aux_dict.get("norms"),
            clipping_norm=current_clipping_norm,
            group_norms_dict=aux_dict.get("group_norms"),
        )

    # Wrap function to accept and return state
    def stateful_clipped_fn(*args, state, **kwargs):
        runtime_clipping_norm = (
            state.clipping_norm if isinstance(state, _RuntimeClipState) else None
        )
        result = clipped_fn(runtime_clipping_norm, *args, **kwargs)
        return result, state  # State unchanged for fixed clipping

    # Return wrapped function with state
    return stateful_clipped_fn, clip_state

`types` ¶

DP-FTRL clipping state and aux types for type annotations.

`ClippedGradResult = ClippedGradValue | tuple[ClippedGradValue, 'ClippedGradAux'] | tuple[ClippedGradValue, 'ClippingStats']` `module-attribute` ¶

Value returned by :class:ClippedGradFn before its updated state.

`ClipPytreeAux = namedtuple('ClipPytreeAux', ['norm', 'group_norms'])` `module-attribute` ¶

Auxiliary outputs from clip_pytree.

Fields

norm: The L2 norm of the original (unclipped) pytree. group_norms: Per-group L2 norms before clipping (dict[str, Tensor]), or None when global clipping is used.

`AutoClippedFunAux` `dataclass` ¶

Bases: ClippedFunAux

Diagnostic outputs from auto_clipped_fun.

Inherits all fields from :class:ClippedFunAux with AUTO-S semantics:

norms: per-example L2 norms before scaling (i.e. ||v||).
clipped_norms: per-example L2 norms after AUTO-S scaling (R * ||v|| / (||v|| + gamma)), clipped by R.
clipping_rate: fraction of examples with ||v|| > R (approximately the fraction where the scale factor was < 1).
values, value_aux, batch_size, group_norms: unchanged from :class:ClippedFunAux.

Source code in packages/opake-engine/src/opake/api/engine/clipping/_auto.py

@dataclass(frozen=True)
class AutoClippedFunAux(ClippedFunAux):
    """Diagnostic outputs from ``auto_clipped_fun``.

    Inherits all fields from :class:`ClippedFunAux` with AUTO-S
    semantics:

    - ``norms``: per-example L2 norms before scaling (i.e. ``||v||``).
    - ``clipped_norms``: per-example L2 norms after AUTO-S scaling
      (``R * ||v|| / (||v|| + gamma)``), clipped by ``R``.
    - ``clipping_rate``: fraction of examples with ``||v|| > R``
      (approximately the fraction where the scale factor was < 1).
    - ``values``, ``value_aux``, ``batch_size``, ``group_norms``:
      unchanged from :class:`ClippedFunAux`.
    """

`AutoClippedGradAux` `dataclass` ¶

Bases: ClippedGradAux

Diagnostic outputs from auto_clipped_grad.

Inherits all fields from :class:ClippedGradAux with AUTO-S semantics:

grad_norms: per-example gradient L2 norms before scaling.
clipped_grad_norms: per-example gradient L2 norms after AUTO-S scaling (R * ||g|| / (||g|| + gamma)), clipped by R.
clipping_rate: fraction of examples with ||g|| > R (approximately the fraction where the scale factor was < 1).
loss_values, loss_aux, batch_size, group_norms: unchanged from :class:ClippedGradAux.

Source code in packages/opake-engine/src/opake/api/engine/clipping/_auto.py

@dataclass(frozen=True)
class AutoClippedGradAux(ClippedGradAux):
    """Diagnostic outputs from ``auto_clipped_grad``.

    Inherits all fields from :class:`ClippedGradAux` with AUTO-S
    semantics:

    - ``grad_norms``: per-example gradient L2 norms before scaling.
    - ``clipped_grad_norms``: per-example gradient L2 norms after
      AUTO-S scaling (``R * ||g|| / (||g|| + gamma)``), clipped by ``R``.
    - ``clipping_rate``: fraction of examples with ``||g|| > R``
      (approximately the fraction where the scale factor was < 1).
    - ``loss_values``, ``loss_aux``, ``batch_size``, ``group_norms``:
      unchanged from :class:`ClippedGradAux`.
    """

`AutoClipState` `dataclass` ¶

Bases: ClipState

Marker state for AUTO-S automatic clipping.

Source code in packages/opake-engine/src/opake/api/engine/clipping/_auto.py

@dataclass(frozen=True)
class AutoClipState(ClipState):
    """Marker state for AUTO-S automatic clipping."""

`ClippedFunAux` `dataclass` ¶

Diagnostic outputs from clipped_fun.

All fields are diagnostic — they reflect pre-noise, pre-aggregation values and must not be fed back into private computation. Use the returned ClippedPytree.max_norm metadata for noise calibration.

Fields

values: Per-example function values before clipping. norms: Per-example L2 norms before clipping. clipped_norms: Per-example L2 norms after clipping. value_aux: Per-example auxiliary payload returned by the wrapped function. clipping_rate: Fraction of per-example outputs whose norm exceeded the clipping threshold. Equal to num_clipped / batch_size. batch_size: Number of examples in the batch. group_norms: Per-group per-example L2 norms before clipping (dict[str, Tensor] with shape [batch_size] per group), or None when global clipping is used.

Source code in packages/opake-engine/src/opake/api/engine/clipping/_clipped_fun.py

@dataclass(frozen=True)
class ClippedFunAux:
    """Diagnostic outputs from clipped_fun.

    All fields are diagnostic — they reflect pre-noise, pre-aggregation
    values and must not be fed back into private computation.  Use the
    returned ``ClippedPytree.max_norm`` metadata for noise calibration.

    Fields:
        values: Per-example function values before clipping.
        norms: Per-example L2 norms before clipping.
        clipped_norms: Per-example L2 norms after clipping.
        value_aux: Per-example auxiliary payload returned by the wrapped function.
        clipping_rate: Fraction of per-example outputs whose norm exceeded the
            clipping threshold.  Equal to ``num_clipped / batch_size``.
        batch_size: Number of examples in the batch.
        group_norms: Per-group per-example L2 norms before clipping
            (dict[str, Tensor] with shape [batch_size] per group), or None
            when global clipping is used.
    """

    values: Any | None = None
    norms: Any | None = None
    clipped_norms: Any | None = None
    value_aux: Any | None = None
    clipping_rate: float | None = None
    batch_size: int = 0
    group_norms: dict[str, torch.Tensor] | None = None

`ClippedGradAux` `dataclass` ¶

Diagnostic outputs from clipped_grad.

All fields are diagnostic — they reflect pre-noise, pre-aggregation values and must not be fed back into private computation. Use the returned ClippedPytree.max_norm metadata for noise calibration.

Fields

loss_values: Per-example loss values before clipping. grad_norms: Per-example gradient L2 norms before clipping. clipped_grad_norms: Per-example gradient L2 norms after clipping. loss_aux: Per-example auxiliary payload returned by the loss function. clipping_rate: Fraction of per-example gradients whose norm exceeded the clipping threshold. batch_size: Number of examples in the batch. group_norms: Per-group per-example gradient L2 norms before clipping (dict[str, Tensor] with shape [batch_size] per group), or None when global clipping is used.

Source code in packages/opake-engine/src/opake/api/engine/clipping/_clipped_grad.py

@dataclass(frozen=True)
class ClippedGradAux:
    """Diagnostic outputs from clipped_grad.

    All fields are diagnostic — they reflect pre-noise, pre-aggregation
    values and must not be fed back into private computation.  Use the
    returned ``ClippedPytree.max_norm`` metadata for noise calibration.

    Fields:
        loss_values: Per-example loss values before clipping.
        grad_norms: Per-example gradient L2 norms before clipping.
        clipped_grad_norms: Per-example gradient L2 norms after clipping.
        loss_aux: Per-example auxiliary payload returned by the loss function.
        clipping_rate: Fraction of per-example gradients whose norm exceeded
            the clipping threshold.
        batch_size: Number of examples in the batch.
        group_norms: Per-group per-example gradient L2 norms before clipping
            (dict[str, Tensor] with shape [batch_size] per group), or None
            when global clipping is used.
    """

    loss_values: Any | None = None
    grad_norms: Any | None = None
    clipped_grad_norms: Any | None = None
    loss_aux: Any | None = None
    clipping_rate: float | None = None
    batch_size: int = 0
    group_norms: dict[str, torch.Tensor] | None = None

`ClippedGradFn` ¶

Bases: Protocol

Callable returned by :func:opake.dpsgd.clipping.clipped_grad.

The positional arguments mirror the wrapped loss function. state is passed by keyword and every invocation returns a new immutable :class:FixedClipState with the clipped gradient result.

Source code in packages/opake-engine/src/opake/api/engine/clipping/_types.py

class ClippedGradFn(Protocol):
    """Callable returned by :func:`opake.dpsgd.clipping.clipped_grad`.

    The positional arguments mirror the wrapped loss function. ``state`` is
    passed by keyword and every invocation returns a new immutable
    :class:`FixedClipState` with the clipped gradient result.
    """

    def __call__(
        self,
        *args: Any,
        state: FixedClipState,
        **kwargs: Any,
    ) -> tuple[ClippedGradResult, FixedClipState]: ...

`FixedClipState` `dataclass` ¶

Bases: ClipState

Marker state for fixed (non-adaptive) clipping.

Returned by :func:clipped_fun and :func:opake.api.engine.clipping.clipped_grad. Carries no fields; the configured clipping threshold flows through the ClippedPytree.max_norm metadata, not through the state.

Source code in packages/opake-engine/src/opake/api/engine/clipping/_clipped_fun.py

@dataclass(frozen=True)
class FixedClipState(_ClipState):
    """Marker state for fixed (non-adaptive) clipping.

    Returned by :func:`clipped_fun` and :func:`opake.api.engine.clipping.clipped_grad`.
    Carries no fields; the configured clipping threshold flows through
    the ``ClippedPytree.max_norm`` metadata, not through the state.
    """

Type	Description
`TypeError`	If `params` has a non-tensor leaf.
`ValueError`	If no patterns are provided, if a parameter matches zero patterns (and no `fallback` is given), if a parameter matches multiple patterns, if an explicit pattern matches no parameter (unless `allow_unused_patterns=True`), or if any value is not positive.

Gradient Clipping¶

Overview¶

Clipping functions¶

State types¶

Auxiliary output types¶

PyTree wrappers¶

Distributed sync helpers¶

API Documentation¶

clipping ¶

auto_clipped_grad(loss_fn, argnums=0, has_aux=False, *, R=1.0, gamma=_DEFAULT_GAMMA, normalize_by=1.0, batch_argnums=1, return_aux=False, return_stats=False, pre_clipping_transform=lambda x: x, microbatch_size=None, dtype=None, second_moment=False, _chunk_compiler=None) ¶

per_group(params, /, patterns=None, *, fallback=None, allow_unused_patterns=False, **kwargs) ¶

fun ¶

auto_clipped_fun(fun, has_aux=False, *, batch_argnums=0, R=1.0, gamma=_DEFAULT_GAMMA, normalize_by=1.0, return_aux=False, return_stats=False, microbatch_size=None, dtype=None, _chunk_compiler=None) ¶

auto_scale_pytree(pytree, R=1.0, gamma=0.01, *, compute_dtype=None) ¶

clip_pytree(pytree, clipping_norm, return_zero=False, *, compute_dtype=None) ¶

clipped_fun(fun, has_aux=False, *, batch_argnums=0, clipping_norm=1.0, normalize_by=1.0, return_aux=False, return_stats=False, second_moment=False, microbatch_size=None, dtype=None, compute_dtype=None, _scale_fn=None, _builtin_scale=None, _chunk_compiler=None) ¶

types ¶

ClippedGradResult = ClippedGradValue | tuple[ClippedGradValue, 'ClippedGradAux'] | tuple[ClippedGradValue, 'ClippingStats'] module-attribute ¶

AdaptiveClippedGradAux dataclass ¶

AdaptiveClipState dataclass ¶

__post_init__() ¶

AutoClippedFunAux dataclass ¶

AutoClippedGradAux dataclass ¶

AutoClipState dataclass ¶

ClippedGradFn ¶

ClippingStats dataclass ¶

FixedClipState dataclass ¶

clipping ¶

auto_clipped_grad(loss_fn, argnums=0, has_aux=False, *, R=1.0, gamma=_DEFAULT_GAMMA, normalize_by=1.0, batch_argnums=1, return_aux=False, return_stats=False, pre_clipping_transform=lambda x: x, microbatch_size=None, dtype=None, second_moment=False, _chunk_compiler=None) ¶

per_group(params, /, patterns=None, *, fallback=None, allow_unused_patterns=False, **kwargs) ¶

fun ¶

auto_clipped_fun(fun, has_aux=False, *, batch_argnums=0, R=1.0, gamma=_DEFAULT_GAMMA, normalize_by=1.0, return_aux=False, return_stats=False, microbatch_size=None, dtype=None, _chunk_compiler=None) ¶

auto_scale_pytree(pytree, R=1.0, gamma=0.01, *, compute_dtype=None) ¶

clip_pytree(pytree, clipping_norm, return_zero=False, *, compute_dtype=None) ¶

clipped_fun(fun, has_aux=False, *, batch_argnums=0, clipping_norm=1.0, normalize_by=1.0, return_aux=False, return_stats=False, second_moment=False, microbatch_size=None, dtype=None, compute_dtype=None, _scale_fn=None, _builtin_scale=None, _chunk_compiler=None) ¶

types ¶

ClippedGradResult = ClippedGradValue | tuple[ClippedGradValue, 'ClippedGradAux'] | tuple[ClippedGradValue, 'ClippingStats'] module-attribute ¶

ClipPytreeAux = namedtuple('ClipPytreeAux', ['norm', 'group_norms']) module-attribute ¶

AutoClippedFunAux dataclass ¶

AutoClippedGradAux dataclass ¶

AutoClipState dataclass ¶

ClippedFunAux dataclass ¶

ClippedGradAux dataclass ¶

ClippedGradFn ¶

FixedClipState dataclass ¶

`clipping` ¶

`auto_clipped_grad(loss_fn, argnums=0, has_aux=False, *, R=1.0, gamma=_DEFAULT_GAMMA, normalize_by=1.0, batch_argnums=1, return_aux=False, return_stats=False, pre_clipping_transform=lambda x: x, microbatch_size=None, dtype=None, second_moment=False, _chunk_compiler=None)` ¶

`per_group(params, /, patterns=None, *, fallback=None, allow_unused_patterns=False, **kwargs)` ¶

`fun` ¶

`auto_clipped_fun(fun, has_aux=False, *, batch_argnums=0, R=1.0, gamma=_DEFAULT_GAMMA, normalize_by=1.0, return_aux=False, return_stats=False, microbatch_size=None, dtype=None, _chunk_compiler=None)` ¶

`auto_scale_pytree(pytree, R=1.0, gamma=0.01, *, compute_dtype=None)` ¶

`clip_pytree(pytree, clipping_norm, return_zero=False, *, compute_dtype=None)` ¶

`clipped_fun(fun, has_aux=False, *, batch_argnums=0, clipping_norm=1.0, normalize_by=1.0, return_aux=False, return_stats=False, second_moment=False, microbatch_size=None, dtype=None, compute_dtype=None, _scale_fn=None, _builtin_scale=None, _chunk_compiler=None)` ¶

`types` ¶

`ClippedGradResult = ClippedGradValue | tuple[ClippedGradValue, 'ClippedGradAux'] | tuple[ClippedGradValue, 'ClippingStats']` `module-attribute` ¶

`AdaptiveClippedGradAux` `dataclass` ¶

`AdaptiveClipState` `dataclass` ¶

`__post_init__()` ¶

`AutoClippedFunAux` `dataclass` ¶

`AutoClippedGradAux` `dataclass` ¶

`AutoClipState` `dataclass` ¶

`ClippedGradFn` ¶

`ClippingStats` `dataclass` ¶

`FixedClipState` `dataclass` ¶

`clipping` ¶

`auto_clipped_grad(loss_fn, argnums=0, has_aux=False, *, R=1.0, gamma=_DEFAULT_GAMMA, normalize_by=1.0, batch_argnums=1, return_aux=False, return_stats=False, pre_clipping_transform=lambda x: x, microbatch_size=None, dtype=None, second_moment=False, _chunk_compiler=None)` ¶

`per_group(params, /, patterns=None, *, fallback=None, allow_unused_patterns=False, **kwargs)` ¶

`fun` ¶

`auto_clipped_fun(fun, has_aux=False, *, batch_argnums=0, R=1.0, gamma=_DEFAULT_GAMMA, normalize_by=1.0, return_aux=False, return_stats=False, microbatch_size=None, dtype=None, _chunk_compiler=None)` ¶

`auto_scale_pytree(pytree, R=1.0, gamma=0.01, *, compute_dtype=None)` ¶

`clip_pytree(pytree, clipping_norm, return_zero=False, *, compute_dtype=None)` ¶

`clipped_fun(fun, has_aux=False, *, batch_argnums=0, clipping_norm=1.0, normalize_by=1.0, return_aux=False, return_stats=False, second_moment=False, microbatch_size=None, dtype=None, compute_dtype=None, _scale_fn=None, _builtin_scale=None, _chunk_compiler=None)` ¶

`types` ¶

`ClippedGradResult = ClippedGradValue | tuple[ClippedGradValue, 'ClippedGradAux'] | tuple[ClippedGradValue, 'ClippingStats']` `module-attribute` ¶

`ClipPytreeAux = namedtuple('ClipPytreeAux', ['norm', 'group_norms'])` `module-attribute` ¶

`AutoClippedFunAux` `dataclass` ¶

`AutoClippedGradAux` `dataclass` ¶

`AutoClipState` `dataclass` ¶

`ClippedFunAux` `dataclass` ¶

`ClippedGradAux` `dataclass` ¶

`ClippedGradFn` ¶

`FixedClipState` `dataclass` ¶