Skip to content

Utilities

Functional Utilities

The opake.functional module provides the torch.func bridges that turn standard nn.Module models into the pure-function form DP-SGD needs. opake.dpsgd.clipping.clipped_grad and friends assume their loss_fn argument is functional in this sense.

Key function: make_functional() — convert a PyTorch module to functional form.

import torch.nn as nn
from opake.functional import make_functional

model = nn.Linear(10, 1)
fmodel, params = make_functional(model)
output = fmodel(params, input_tensor)

Functional models allow torch.func.vmap to compute per-example gradients efficiently, which is essential for DP-SGD.

Trainable / frozen partition (for PEFT and LoRA)

For parameter-efficient fine-tuning (LoRA, adapters, BitFit, …) only a small subset of parameters has requires_grad=True. Pass partition_trainable=True to split the parameter dict accordingly:

fmodel, trainable, frozen = make_functional(model, partition_trainable=True)

def loss_fn(trainable_params, x, y):
    out = fmodel({**frozen, **trainable_params}, x)
    return ((out - y) ** 2).mean()

Only trainable_params flows into the loss closure, so vmap(grad(...)) over loss_fn produces per-example gradients only for the trainable subset. Frozen parameters are broadcast (constant) under vmap, which is what makes LoRA-style DP fine-tuning of multi-billion-parameter models feasible — per-example gradient memory scales with trainable params, not total params. See Memory optimizations for the memory arithmetic; fused LoRA Triton kernels are documented under Model patches — Fused LoRA operations.

See also: Quick Start Guide for functional API usage.

functional

Functional API — make_functional, with_batch_dim, empty_collate.

empty_collate(collate_fn)

Wrap a collate function so empty example lists yield zero-batch outputs.

Poisson-style samplers can yield empty index lists, producing batches with zero examples. Standard collators (e.g., HuggingFace DataCollatorForLanguageModeling) crash on empty input because they unconditionally index examples[0].

This wrapper learns the output structure from the first non-empty batch and returns a properly-shaped empty batch (batch dim = 0) when the example list is empty::

collate = empty_collate(my_collate_fn)
loader = DataLoader(dataset, batch_sampler=sampler, collate_fn=collate)

for batch in loader:
    input_ids = batch["input_ids"]   # shape (0, seq_len) when empty
    if len(input_ids) == 0:
        continue
    ...

.. note::

If the very first batch is empty (before any structure has been
learned), the call falls through to *collate_fn* which will
typically crash — this is the correct signal that the sample rate
is misconfigured.

Parameters:

Name Type Description Default
collate_fn Callable[..., T]

Original collate function.

 required

Returns:

Type Description
Callable[..., T]

Wrapped function that returns empty-batch-dim outputs for empty
Callable[..., T]

example lists.
Source code in packages/opake-engine/src/opake/api/engine/functional/_collate.py 
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
def empty_collate(collate_fn: Callable[..., T]) -> Callable[..., T]:
    """Wrap a collate function so empty example lists yield zero-batch outputs.

    Poisson-style samplers can yield empty index lists, producing batches with
    zero examples. Standard collators (e.g., HuggingFace
    ``DataCollatorForLanguageModeling``) crash on empty input because they
    unconditionally index ``examples[0]``.

    This wrapper learns the output structure from the first non-empty batch
    and returns a properly-shaped empty batch (batch dim = 0) when the
    example list is empty::

        collate = empty_collate(my_collate_fn)
        loader = DataLoader(dataset, batch_sampler=sampler, collate_fn=collate)

        for batch in loader:
            input_ids = batch["input_ids"]   # shape (0, seq_len) when empty
            if len(input_ids) == 0:
                continue
            ...

    .. note::

        If the very first batch is empty (before any structure has been
        learned), the call falls through to *collate_fn* which will
        typically crash — this is the correct signal that the sample rate
        is misconfigured.

    Args:
        collate_fn: Original collate function.

    Returns:
        Wrapped function that returns empty-batch-dim outputs for empty
        example lists.
    """
    return _EmptyCollator(collate_fn)

make_functional(mod, disable_autograd_tracking=False, partition_trainable=False)

Convert a PyTorch module to functional form.

This helper mimics the behavior of the deprecated functorch.make_functional(). It creates a stateless version of the module that can be called with external parameters, enabling use with torch.func transformations like vmap and grad.

The returned functional model can be used with per-example gradient computation, which is essential for differential privacy.

Parameters:

Name Type Description Default
mod Module

PyTorch module to convert to functional form.

 required
disable_autograd_tracking bool

If True, detach parameters from autograd graph. Useful when you only need to compute gradients via torch.func.grad, not through standard .backward(). Default: False.

 False
partition_trainable bool

If True, return (fmodel, trainable_params, frozen_params) as dicts instead of (fmodel, params) as tuple. Partitioning is based on the requires_grad attribute of the original module's parameters. This is ideal for LoRA-style fine-tuning. Default: False.

 False

Returns:

Type Description
tuple[Callable, tuple[Tensor, ...]] | tuple[Callable, dict, dict]

If partition_trainable=False: A tuple containing: - fmodel: Functional version of the module. Takes parameters as first argument, followed by the module's normal forward arguments. - params: Tuple of the module's parameter tensors.
tuple[Callable, tuple[Tensor, ...]] | tuple[Callable, dict, dict]

If partition_trainable=True: A tuple containing: - fmodel: Functional version of the module. Takes parameters dict as first argument. - trainable_params: Dict of parameters where requires_grad=True. - frozen_params: Dict of parameters where requires_grad=False.
Example

import torch import torch.nn as nn from opake.api.engine.functional_utils import make_functional

Create a simple model

model = nn.Linear(10, 1)

Convert to functional form

fmodel, params = make_functional(model)

Use with new parameters

x = torch.randn(5, 10) output = fmodel(params, x) print(output.shape) torch.Size([5, 1])

Works with torch.func transformations

from torch.func import grad, vmap

def loss_fn(params, x, y): ... pred = fmodel(params, x) ... return ((pred - y) ** 2).mean()

Compute gradient w.r.t. parameters

x_batch = torch.randn(3, 10) y_batch = torch.randn(3, 1) grads = grad(loss_fn)(params, x_batch, y_batch)

Example with per-example gradients
Per-example loss function

def loss_single(params, x, y): ... pred = fmodel(params, x.unsqueeze(0)) ... return ((pred - y) ** 2).mean()

Compute per-example gradients with vmap

per_example_grads = vmap(grad(loss_single), in_dims=(None, 0, 0))( ... params, x_batch, y_batch ... )

per_example_grads is a tuple with shape (batch_size, *param_shape) for each param
Example with LoRA partitioning

import torch.nn as nn from opake.dpsgd.clipping import clipped_grad from opake.dpsgd.noise import gaussian_noise from opake.api.engine.random import key from opake.api.engine.functional import make_functional, merge

Create model and freeze backbone

model = nn.Sequential( ... nn.Linear(10, 5, bias=False), # Pretrained backbone ... nn.Linear(5, 1, bias=False), # Pretrained backbone ... )

Freeze first layer, keep second trainable

model[0].weight.requires_grad = False # Frozen model[1].weight.requires_grad = True # Trainable

Convert with automatic partitioning

fmodel, trainable, frozen = make_functional( ... model, partition_trainable=True ... )

Training loop

def loss_fn(train_params, x, y): ... all_params = merge(frozen, train_params) ... pred = fmodel(all_params, x) ... return ((pred - y) ** 2).mean()

grad_fn, clip_state = clipped_grad( ... loss_fn, argnums=0, batch_argnums=(1, 2), clipping_norm=1.0 ... ) noise_fn, noise_state = gaussian_noise(noise_multiplier=1.1, key=key(42)) optimizer = torch.optim.SGD([p for p in model.parameters() if p.requires_grad], lr=1e-3)

for x_batch, y_batch in dataloader: ... grads, clip_state = grad_fn(trainable, x_batch, y_batch, state=clip_state) ... noisy_grads, noise_state = noise_fn(grads, noise_state) ... # ... assign grads and step optimizer

Note

Gradient checkpointing compatibility is handled by the reparametrize_recompute patch in opake.api.patches.torch.checkpoint, which re-binds the functional_call parameters during checkpoint recomputation in backward (on a PyTorch with native support the patch is skipped and torch handles this itself). This wrapper is purely functional — it delegates entirely to torch.func.functional_call.

See Also
  • PyTorch migration guide: https://pytorch.org/docs/master/func.migrating.html
Source code in packages/opake-engine/src/opake/api/engine/functional/__init__.py 
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
def make_functional(
    mod: nn.Module,
    disable_autograd_tracking: bool = False,
    partition_trainable: bool = False,
) -> tuple[Callable, tuple[torch.Tensor, ...]] | tuple[Callable, dict, dict]:
    """Convert a PyTorch module to functional form.

    This helper mimics the behavior of the deprecated `functorch.make_functional()`.
    It creates a stateless version of the module that can be called with external
    parameters, enabling use with `torch.func` transformations like `vmap` and `grad`.

    The returned functional model can be used with per-example gradient computation,
    which is essential for differential privacy.

    Args:
        mod: PyTorch module to convert to functional form.
        disable_autograd_tracking: If True, detach parameters from autograd graph.
            Useful when you only need to compute gradients via torch.func.grad,
            not through standard .backward(). Default: False.
        partition_trainable: If True, return (fmodel, trainable_params, frozen_params)
            as dicts instead of (fmodel, params) as tuple. Partitioning is based on
            the `requires_grad` attribute of the original module's parameters.
            This is ideal for LoRA-style fine-tuning. Default: False.

    Returns:
        If partition_trainable=False:
            A tuple containing:
            - fmodel: Functional version of the module. Takes parameters as first argument,
              followed by the module's normal forward arguments.
            - params: Tuple of the module's parameter tensors.

        If partition_trainable=True:
            A tuple containing:
            - fmodel: Functional version of the module. Takes parameters dict as first argument.
            - trainable_params: Dict of parameters where requires_grad=True.
            - frozen_params: Dict of parameters where requires_grad=False.

    Example:
        >>> import torch
        >>> import torch.nn as nn
        >>> from opake.api.engine.functional_utils import make_functional
        >>>
        >>> # Create a simple model
        >>> model = nn.Linear(10, 1)
        >>>
        >>> # Convert to functional form
        >>> fmodel, params = make_functional(model)
        >>>
        >>> # Use with new parameters
        >>> x = torch.randn(5, 10)
        >>> output = fmodel(params, x)
        >>> print(output.shape)
        torch.Size([5, 1])
        >>>
        >>> # Works with torch.func transformations
        >>> from torch.func import grad, vmap
        >>>
        >>> def loss_fn(params, x, y):
        ...     pred = fmodel(params, x)
        ...     return ((pred - y) ** 2).mean()
        >>>
        >>> # Compute gradient w.r.t. parameters
        >>> x_batch = torch.randn(3, 10)
        >>> y_batch = torch.randn(3, 1)
        >>> grads = grad(loss_fn)(params, x_batch, y_batch)

    Example with per-example gradients:
        >>> # Per-example loss function
        >>> def loss_single(params, x, y):
        ...     pred = fmodel(params, x.unsqueeze(0))
        ...     return ((pred - y) ** 2).mean()
        >>>
        >>> # Compute per-example gradients with vmap
        >>> per_example_grads = vmap(grad(loss_single), in_dims=(None, 0, 0))(
        ...     params, x_batch, y_batch
        ... )
        >>> # per_example_grads is a tuple with shape (batch_size, *param_shape) for each param

    Example with LoRA partitioning:
        >>> import torch.nn as nn
        >>> from opake.dpsgd.clipping import clipped_grad
        >>> from opake.dpsgd.noise import gaussian_noise
        >>> from opake.api.engine.random import key
        >>> from opake.api.engine.functional import make_functional, merge
        >>>
        >>> # Create model and freeze backbone
        >>> model = nn.Sequential(
        ...     nn.Linear(10, 5, bias=False),  # Pretrained backbone
        ...     nn.Linear(5, 1, bias=False),   # Pretrained backbone
        ... )
        >>>
        >>> # Freeze first layer, keep second trainable
        >>> model[0].weight.requires_grad = False  # Frozen
        >>> model[1].weight.requires_grad = True   # Trainable
        >>>
        >>> # Convert with automatic partitioning
        >>> fmodel, trainable, frozen = make_functional(
        ...     model, partition_trainable=True
        ... )
        >>>
        >>> # Training loop
        >>> def loss_fn(train_params, x, y):
        ...     all_params = merge(frozen, train_params)
        ...     pred = fmodel(all_params, x)
        ...     return ((pred - y) ** 2).mean()
        >>>
        >>> grad_fn, clip_state = clipped_grad(
        ...     loss_fn, argnums=0, batch_argnums=(1, 2), clipping_norm=1.0
        ... )
        >>> noise_fn, noise_state = gaussian_noise(noise_multiplier=1.1, key=key(42))
        >>> optimizer = torch.optim.SGD([p for p in model.parameters() if p.requires_grad], lr=1e-3)
        >>>
        >>> for x_batch, y_batch in dataloader:
        ...     grads, clip_state = grad_fn(trainable, x_batch, y_batch, state=clip_state)
        ...     noisy_grads, noise_state = noise_fn(grads, noise_state)
        ...     # ... assign grads and step optimizer

    Note:
        Gradient checkpointing compatibility is handled by the
        ``reparametrize_recompute`` patch in
        ``opake.api.patches.torch.checkpoint``, which re-binds the
        ``functional_call`` parameters during checkpoint recomputation in
        backward (on a PyTorch with native support the patch is skipped and
        torch handles this itself).  This wrapper is purely functional — it
        delegates entirely to ``torch.func.functional_call``.

    See Also:
        - PyTorch migration guide: https://pytorch.org/docs/master/func.migrating.html
    """
    # Extract parameters with their requires_grad status
    params_dict = dict(mod.named_parameters())

    stateless_mod = mod

    # Optionally detach parameters from autograd graph
    if disable_autograd_tracking:
        params_dict = {name: param.detach() for name, param in params_dict.items()}

    if partition_trainable:
        # Partition based on requires_grad (BEFORE detaching!)
        # We need to check the original parameters
        original_params = dict(mod.named_parameters())
        trainable_params = {
            name: params_dict[name]
            for name, param in original_params.items()
            if param.requires_grad
        }
        frozen_params = {
            name: params_dict[name]
            for name, param in original_params.items()
            if not param.requires_grad
        }

        def fmodel_dict(params_dict_input, *args, **kwargs):
            return torch.func.functional_call(
                stateless_mod, params_dict_input, args, kwargs
            )

        return fmodel_dict, trainable_params, frozen_params

    else:
        # Original behavior: return tuple
        params_names = list(params_dict.keys())
        params_values = tuple(params_dict.values())

        def fmodel_tuple(new_params_values, *args, **kwargs):
            new_params_dict = dict(zip(params_names, new_params_values, strict=True))
            return torch.func.functional_call(
                stateless_mod, new_params_dict, args, kwargs
            )

        return fmodel_tuple, params_values

with_batch_dim(fn, batch_argnums=(), batch_kwargs=(), *, min_ndim=None)

Wrap a function to add a leading batch dimension to specified arguments.

Use this when a model expects a batch dimension but the function is called per-example under vmap where the batch dim has been stripped.

Two modes controlled by min_ndim:

  • min_ndim=None (default): always unsqueeze specified args, never touch output. This is the user-facing mode for wrapping loss functions.

  • min_ndim=N (idempotent mode): only unsqueeze tensors with ndim < threshold, and squeeze top-level output tensors when any unsqueezing happened. This is the model-forward mode — a no-op for already-batched inputs.

Parameters:

Name Type Description Default
fn Callable

Function to wrap.

 required
batch_argnums int | tuple[int, ...]

Which positional arguments should get an extra leading dimension via unsqueeze(0). Int or tuple of ints.

 ()
batch_kwargs tuple[str, ...] | dict[str, int]

Keyword arguments to process. A tuple of names uses the global min_ndim threshold. A dict maps each kwarg name to its own ndim threshold (e.g., {"input_ids": 2, "inputs_embeds": 3}).

 ()
min_ndim int | None

Global ndim threshold. None means always unsqueeze. An int means only unsqueeze when tensor.ndim < min_ndim.

 None

Returns:

Type Description
Callable

Wrapped function.

Example — loss function mode (always unsqueeze)::

>>> def loss_fn(params, tokens):
...     output = model(params, tokens)  # model expects (batch, seq)
...     return output.loss
>>> # Under vmap, tokens is (seq,); wrap so model sees (1, seq)
>>> loss_fn = with_batch_dim(loss_fn, batch_argnums=1)

Example — model forward mode (idempotent)::

>>> # Only unsqueeze when inputs lack batch dim; squeeze output
>>> model.forward = with_batch_dim(
...     model.forward,
...     batch_kwargs={"input_ids": 2, "inputs_embeds": 3},
...     min_ndim=2,
... )
Source code in packages/opake-engine/src/opake/api/engine/functional/__init__.py 
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
def with_batch_dim(
    fn: Callable,
    batch_argnums: int | tuple[int, ...] = (),
    batch_kwargs: tuple[str, ...] | dict[str, int] = (),
    *,
    min_ndim: int | None = None,
) -> Callable:
    """Wrap a function to add a leading batch dimension to specified arguments.

    Use this when a model expects a batch dimension but the function is called
    per-example under ``vmap`` where the batch dim has been stripped.

    Two modes controlled by ``min_ndim``:

    - ``min_ndim=None`` (default): always unsqueeze specified args, never
      touch output. This is the user-facing mode for wrapping loss functions.

    - ``min_ndim=N`` (idempotent mode): only unsqueeze tensors with
      ``ndim < threshold``, and squeeze top-level output tensors when any
      unsqueezing happened. This is the model-forward mode — a no-op for
      already-batched inputs.

    Args:
        fn: Function to wrap.
        batch_argnums: Which positional arguments should get an extra leading
            dimension via ``unsqueeze(0)``. Int or tuple of ints.
        batch_kwargs: Keyword arguments to process. A tuple of names uses the
            global ``min_ndim`` threshold. A dict maps each kwarg name to its
            own ndim threshold (e.g., ``{"input_ids": 2, "inputs_embeds": 3}``).
        min_ndim: Global ndim threshold. ``None`` means always unsqueeze.
            An int means only unsqueeze when ``tensor.ndim < min_ndim``.

    Returns:
        Wrapped function.

    Example — loss function mode (always unsqueeze)::

        >>> def loss_fn(params, tokens):
        ...     output = model(params, tokens)  # model expects (batch, seq)
        ...     return output.loss
        >>> # Under vmap, tokens is (seq,); wrap so model sees (1, seq)
        >>> loss_fn = with_batch_dim(loss_fn, batch_argnums=1)

    Example — model forward mode (idempotent)::

        >>> # Only unsqueeze when inputs lack batch dim; squeeze output
        >>> model.forward = with_batch_dim(
        ...     model.forward,
        ...     batch_kwargs={"input_ids": 2, "inputs_embeds": 3},
        ...     min_ndim=2,
        ... )
    """
    if getattr(fn, "_opake_batchified", False):
        return fn

    import inspect

    from opake.api.engine.pytree import tree_map

    if isinstance(batch_argnums, int):
        batch_argnums = (batch_argnums,)

    # Normalize batch_kwargs: tuple → dict using global min_ndim
    if isinstance(batch_kwargs, tuple):
        batch_kwargs_dict: dict[str, int | None] = dict.fromkeys(batch_kwargs, min_ndim)
    else:
        batch_kwargs_dict = dict(batch_kwargs)

    conditional = min_ndim is not None

    # Pre-compute: positional indices of batch_kwargs parameters in fn's
    # signature, so we can catch them even when passed positionally.
    _positional_kwarg_indices: dict[int, str] = {}
    if batch_kwargs_dict:
        try:
            sig = inspect.signature(fn)
            for idx, (name, param) in enumerate(sig.parameters.items()):
                if param.kind in (param.POSITIONAL_ONLY, param.POSITIONAL_OR_KEYWORD):
                    if name in batch_kwargs_dict:
                        _positional_kwarg_indices[idx] = name
                elif param.kind == param.VAR_POSITIONAL:
                    break
        except (ValueError, TypeError):
            pass

    def _needs_unsqueeze(tensor: torch.Tensor, threshold: int | None) -> bool:
        if threshold is None:
            return True
        return tensor.ndim < threshold

    # ``functools.wraps`` propagates ``fn.__wrapped__`` (and sets it to
    # ``fn`` itself) so callers can walk the chain back to the original
    # function via ``inspect.signature(wrapper, follow_wrapped=True)`` — the
    # default behavior of :func:`inspect.signature`.  Concretely: when this
    # wraps a transformers ``forward(input_ids=None, attention_mask=None,
    # …)``, ``inspect.signature`` resolves to the real named parameters
    # instead of just ``(*args, **kwargs)``.  HF-style column pruning that
    # introspects ``model.forward`` (and any other downstream consumer of
    # the signature) keeps working after batchifying.
    @functools.wraps(fn)
    def wrapper(*args, **kwargs):
        unsqueezed = False
        args_list = list(args)

        # Move batch_kwargs that were passed positionally into kwargs.
        for idx in sorted(_positional_kwarg_indices, reverse=True):
            if idx < len(args_list):
                name = _positional_kwarg_indices[idx]
                if name not in kwargs:
                    kwargs[name] = args_list.pop(idx)

        # Process positional args
        for i in batch_argnums:

            def _unsqueeze_arg(x, _threshold=min_ndim):
                nonlocal unsqueezed
                if isinstance(x, torch.Tensor) and _needs_unsqueeze(x, _threshold):
                    unsqueezed = True
                    return x.unsqueeze(0)
                return x

            if i < len(args_list):
                args_list[i] = tree_map(_unsqueeze_arg, args_list[i])

        # Process keyword args
        for name, threshold in batch_kwargs_dict.items():
            if name not in kwargs or kwargs[name] is None:
                continue
            val = kwargs[name]
            if isinstance(val, torch.Tensor) and _needs_unsqueeze(val, threshold):
                unsqueezed = True
                kwargs[name] = val.unsqueeze(0)

        result = fn(*args_list, **kwargs)

        if conditional and unsqueezed:
            result = _squeeze_output(result)

        return result

    wrapper._opake_batchified = True
    return wrapper

PyTree Utilities

The opake.pytree module provides helpers for working with PyTrees — nested structures of tensors used throughout Opake.

PyTrees are nested dictionaries or tuples of tensors, commonly used to represent model parameters:

params = {
    "encoder": {"weight": tensor1, "bias": tensor2},
    "decoder": {"weight": tensor3, "bias": tensor4},
}

This module provides:

  • tree_map() — apply a function to all leaves
  • tree_map_with_path() — apply a function with key paths to all leaves
  • tree_leaves() — extract all leaf tensors
  • partition() — split a tree into two by predicate
  • merge() — recombine partitioned trees
  • global_norm() — compute L2 norm across the entire tree

See also: Gradient Clipping Guide for PyTree usage in DP-SGD.

pytree

Torch-pytree ops — tree_map, tree_leaves, partition, merge, global_norm.

ParamPath = tuple[str | int, ...] module-attribute

Optree leaf path: nested dict keys (str) and sequence indices (int).

Flat named_parameters dicts use a one-segment path, e.g. ("layers.0.weight",). Nested trees use multi-segment paths, e.g. ("layers", 0, "weight"). The two never collide.

A bare leaf pytree (e.g. a single Tensor) uses the empty path (), matching :func:optree.tree_flatten_with_path.

global_norm(tree, *, compute_dtype=None)

Compute global L2 norm across all tensors in a PyTree.

The global norm is the square root of the sum of squared norms of all leaf tensors: global_norm = sqrt(sum(||leaf||^2 for leaf in tree))

Parameters:

Name Type Description Default
tree Any

PyTree of tensors (e.g., parameters or gradients)

 required
compute_dtype dtype | None

Internal accumulation dtype. None (default) promotes low-precision floats (fp16/bf16) to float32 for numerical stability and otherwise uses the input float dtype promoted to at least float32. Pass an explicit dtype (e.g. torch.float64) to force a specific accumulation precision.

 None

Returns:

Type Description
Tensor

Scalar tensor containing the global L2 norm on the device of the first
Tensor

tensor leaf (or CPU if the tree is empty). Output dtype matches
Tensor

the resolved compute dtype.

Raises:

Type Description
InputTypeError

If tree contains a DP wrapper or compute_dtype is not a real floating-point dtype.
Example

tree = {'w': torch.tensor([3.0, 4.0]), 'b': torch.tensor([0.0, 12.0])} norm = global_norm(tree)

norm = sqrt(3^2 + 4^2 + 0^2 + 12^2) = sqrt(169) = 13

torch.isclose(norm, torch.tensor(13.0)) True

References

This function is commonly used in gradient clipping for deep learning. See: Pascanu et al. 2013, "On the difficulty of training RNNs"

Source code in packages/opake-engine/src/opake/api/engine/pytree.py 
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
def global_norm(
    tree: Any,
    *,
    compute_dtype: torch.dtype | None = None,
) -> torch.Tensor:
    """Compute global L2 norm across all tensors in a PyTree.

    The global norm is the square root of the sum of squared norms of all
    leaf tensors:
        global_norm = sqrt(sum(||leaf||^2 for leaf in tree))

    Args:
        tree: PyTree of tensors (e.g., parameters or gradients)
        compute_dtype: Internal accumulation dtype.  ``None`` (default)
            promotes low-precision floats (fp16/bf16) to float32 for
            numerical stability and otherwise uses the input float dtype
            promoted to at least float32.  Pass an explicit dtype (e.g.
            ``torch.float64``) to force a specific accumulation precision.

    Returns:
        Scalar tensor containing the global L2 norm on the device of the first
        tensor leaf (or CPU if the tree is empty).  Output dtype matches
        the resolved compute dtype.

    Raises:
        InputTypeError: If ``tree`` contains a DP wrapper or ``compute_dtype``
            is not a real floating-point dtype.

    Example:
        >>> tree = {'w': torch.tensor([3.0, 4.0]), 'b': torch.tensor([0.0, 12.0])}
        >>> norm = global_norm(tree)
        >>> # norm = sqrt(3^2 + 4^2 + 0^2 + 12^2) = sqrt(169) = 13
        >>> torch.isclose(norm, torch.tensor(13.0))
        True

    References:
        This function is commonly used in gradient clipping for deep learning.
        See: Pascanu et al. 2013, "On the difficulty of training RNNs"
    """
    from opake.api.engine.types import (
        ClippedPytree,
        NoisedPytree,
        SecondMomentClippingOutput,
        SecondMomentNoiseOutput,
    )

    if isinstance(
        tree,
        (
            ClippedPytree,
            NoisedPytree,
            SecondMomentClippingOutput,
            SecondMomentNoiseOutput,
        ),
    ):
        raise InputTypeError(_RAW_TENSOR_PYTREE_ERROR)

    if compute_dtype is not None and not torch.is_floating_point(
        torch.empty((), dtype=compute_dtype)
    ):
        raise InputTypeError(
            *(
                f"compute_dtype must be a real floating-point dtype, got "
                f"{compute_dtype!r}.  Integer/bool/complex compute dtypes can "
                f"silently corrupt the L2-norm reduction (the squared sum is "
                f"non-negative real and the final sqrt assumes a real "
                f"accumulator).",
            )
        )

    leaves = tree_leaves(tree)
    if not leaves:
        return torch.tensor(0.0, dtype=compute_dtype or torch.float32)

    if compute_dtype is not None:
        real_acc_dtype = compute_dtype
    else:
        # Auto-promote: at least float32, but match user's intent if they
        # supplied higher-precision inputs.
        dtypes = [t.dtype for t in leaves]
        float_dtypes = [
            dt for dt in dtypes if torch.is_floating_point(torch.empty((), dtype=dt))
        ]
        complex_dtypes = [
            dt for dt in dtypes if torch.is_complex(torch.empty((), dtype=dt))
        ]

        if complex_dtypes:
            acc_dtype = torch.complex64
            for dt in complex_dtypes:
                acc_dtype = torch.promote_types(acc_dtype, dt)
            # For complex, accumulate real magnitudes in the corresponding real dtype
            real_acc_dtype = torch.promote_types(
                torch.float32, torch.tensor(0, dtype=acc_dtype).real.dtype
            )
        elif float_dtypes:
            real_acc_dtype = torch.float32
            for dt in float_dtypes:
                real_acc_dtype = torch.promote_types(real_acc_dtype, dt)
        else:
            # All integer/bool → accumulate in float32
            real_acc_dtype = torch.float32

    device = leaves[0].device
    total = torch.zeros((), device=device, dtype=real_acc_dtype)

    for leaf in leaves:
        if torch.is_complex(leaf):
            # ||z||^2 = (real^2 + imag^2)
            real = leaf.real.to(real_acc_dtype)
            imag = leaf.imag.to(real_acc_dtype)
            total = (
                total
                + (real * real).sum(dtype=real_acc_dtype)
                + (imag * imag).sum(dtype=real_acc_dtype)
            )
        else:
            x = leaf.to(real_acc_dtype)
            total = total + (x * x).sum(dtype=real_acc_dtype)

    return torch.sqrt(total)

merge(*trees)

Merge multiple PyTrees into one.

Later trees override earlier ones for overlapping keys.

Parameters:

Name Type Description Default
*trees Any

PyTrees to merge (must have compatible structures)

 ()

Returns:

Type Description
Any

Merged PyTree
Example

trainable = {'encoder': {'lora_a': torch.ones(2)}} frozen = {'encoder': {'weight': torch.zeros(3)}} merged = merge(frozen, trainable) merged {'encoder': {'weight': ..., 'lora_a': ...}}

References

Inspired by Haiku's hk.data_structures.merge(): https://github.com/deepmind/dm-haiku/blob/main/haiku/_src/data_structures.py

Source code in packages/opake-engine/src/opake/api/engine/pytree.py 
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
def merge(*trees: Any) -> Any:
    """Merge multiple PyTrees into one.

    Later trees override earlier ones for overlapping keys.

    Args:
        *trees: PyTrees to merge (must have compatible structures)

    Returns:
        Merged PyTree

    Example:
        >>> trainable = {'encoder': {'lora_a': torch.ones(2)}}
        >>> frozen = {'encoder': {'weight': torch.zeros(3)}}
        >>> merged = merge(frozen, trainable)
        >>> merged
        {'encoder': {'weight': ..., 'lora_a': ...}}

    References:
        Inspired by Haiku's hk.data_structures.merge():
        https://github.com/deepmind/dm-haiku/blob/main/haiku/_src/data_structures.py
    """
    if not trees:
        return {}

    # Filter out None/empty trees
    trees = [t for t in trees if t is not None]
    if not trees:
        return {}

    # Base case: single tree
    if len(trees) == 1:
        return trees[0]

    # Get first tree as reference
    result = trees[0]

    # Merge remaining trees
    for tree in trees[1:]:
        result = _merge_two(result, tree)

    return result

param_path(path)

Normalize an optree path (or flat string key) to :data:ParamPath.

The empty path () is valid and denotes a root leaf (a pytree that is itself a single tensor).

Source code in packages/opake-engine/src/opake/api/engine/pytree.py 
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
def param_path(path: tuple[Any, ...] | list[Any] | str) -> ParamPath:
    """Normalize an optree path (or flat string key) to :data:`ParamPath`.

    The empty path ``()`` is valid and denotes a root leaf (a pytree that
    is itself a single tensor).
    """
    if isinstance(path, str):
        return (path,)
    out: list[str | int] = []
    for part in path:
        if isinstance(part, (str, int)):
            out.append(part)
        else:
            raise InputTypeError(
                *(
                    f"ParamPath components must be str or int; got {type(part).__name__}",
                )
            )
    return tuple(out)

param_path_display(path)

Dotted display form for pattern matching and diagnostics.

Source code in packages/opake-engine/src/opake/api/engine/pytree.py 
72
73
74
def param_path_display(path: ParamPath) -> str:
    """Dotted display form for pattern matching and diagnostics."""
    return ".".join(str(p) for p in path)

partition(predicate, tree)

Partition PyTree into two trees based on predicate.

This is the key function for LoRA-style training: split parameters into trainable (e.g., LoRA adapters) and frozen (e.g., pretrained backbone).

Parameters:

Name Type Description Default
predicate Callable[[tuple[Any, ...], Any], bool]

Function(path, value) -> bool. Returns True for first tree.

 required
tree Any

PyTree to partition

 required

Returns:

Type Description
Any

Tuple of (matching_tree, non_matching_tree) with same structure as input.
Any

Branches where all leaves are filtered out are omitted.
Example

params = { ... 'encoder': { ... 'weight': torch.randn(10, 5), ... 'lora_a': torch.randn(10, 2), ... 'lora_b': torch.randn(2, 5), ... } ... } def is_lora(path, value): ... return 'lora' in str(path) trainable, frozen = partition(is_lora, params) 'lora_a' in trainable['encoder'] # True 'weight' in frozen['encoder'] # True 'weight' in trainable['encoder'] # False

References

Inspired by Haiku's hk.data_structures.partition(): https://github.com/deepmind/dm-haiku/blob/main/haiku/_src/data_structures.py

Source code in packages/opake-engine/src/opake/api/engine/pytree.py 
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
def partition(
    predicate: Callable[[tuple[Any, ...], Any], bool],
    tree: Any,
) -> tuple[Any, Any]:
    """Partition PyTree into two trees based on predicate.

    This is the key function for LoRA-style training: split parameters into
    trainable (e.g., LoRA adapters) and frozen (e.g., pretrained backbone).

    Args:
        predicate: Function(path, value) -> bool. Returns True for first tree.
        tree: PyTree to partition

    Returns:
        Tuple of (matching_tree, non_matching_tree) with same structure as input.
        Branches where all leaves are filtered out are omitted.

    Example:
        >>> params = {
        ...     'encoder': {
        ...         'weight': torch.randn(10, 5),
        ...         'lora_a': torch.randn(10, 2),
        ...         'lora_b': torch.randn(2, 5),
        ...     }
        ... }
        >>> def is_lora(path, value):
        ...     return 'lora' in str(path)
        >>> trainable, frozen = partition(is_lora, params)
        >>> 'lora_a' in trainable['encoder']  # True
        >>> 'weight' in frozen['encoder']      # True
        >>> 'weight' in trainable['encoder']   # False

    References:
        Inspired by Haiku's hk.data_structures.partition():
        https://github.com/deepmind/dm-haiku/blob/main/haiku/_src/data_structures.py
    """

    def _partition_subtree(path: tuple, subtree: Any) -> tuple[Any, Any]:
        if isinstance(subtree, dict):
            true_dict = {}
            false_dict = {}
            for k, v in subtree.items():
                true_part, false_part = _partition_subtree((*path, k), v)
                if true_part is not None:
                    true_dict[k] = true_part
                if false_part is not None:
                    false_dict[k] = false_part
            return (
                true_dict or None,
                false_dict or None,
            )
        elif isinstance(subtree, (list, tuple)):
            true_list = []
            false_list = []
            for i, v in enumerate(subtree):
                true_part, false_part = _partition_subtree((*path, i), v)
                true_list.append(true_part)
                false_list.append(false_part)
            # Keep same type (list vs tuple)
            any_true = any(x is not None for x in true_list)
            any_false = any(x is not None for x in false_list)
            return (
                type(subtree)(true_list) if any_true else None,
                type(subtree)(false_list) if any_false else None,
            )
        else:
            # Leaf node - apply predicate
            if predicate(path, subtree):
                return subtree, None
            else:
                return None, subtree

    true_tree, false_tree = _partition_subtree((), tree)
    # Return empty dicts instead of None for empty trees
    if true_tree is None:
        true_tree = {} if isinstance(tree, dict) else type(tree)()
    if false_tree is None:
        false_tree = {} if isinstance(tree, dict) else type(tree)()
    return true_tree, false_tree

tree_flatten(tree)

Flatten tree to (leaves, treedef).

Thin wrapper around :func:optree.tree_flatten. Prefer :func:tree_flatten_with_paths when callers need leaf identities.

Source code in packages/opake-engine/src/opake/api/engine/pytree.py 
90
91
92
93
94
95
96
97
def tree_flatten(tree: Any) -> tuple[list[Any], Any]:
    """Flatten ``tree`` to ``(leaves, treedef)``.

    Thin wrapper around :func:`optree.tree_flatten`.  Prefer
    :func:`tree_flatten_with_paths` when callers need leaf identities.
    """
    leaves, treedef = _ot.tree_flatten(tree)
    return list(leaves), treedef

tree_flatten_with_paths(tree)

Flatten tree to (paths, leaves, treedef) with :data:ParamPath keys.

Wraps :func:optree.tree_flatten_with_path so PerGroup consumers share one path convention.

Source code in packages/opake-engine/src/opake/api/engine/pytree.py 
77
78
79
80
81
82
83
84
85
86
87
def tree_flatten_with_paths(
    tree: Any,
) -> tuple[list[ParamPath], list[Any], Any]:
    """Flatten ``tree`` to ``(paths, leaves, treedef)`` with :data:`ParamPath` keys.

    Wraps :func:`optree.tree_flatten_with_path` so PerGroup consumers share one
    path convention.
    """
    raw_paths, leaves, treedef = _ot.tree_flatten_with_path(tree)
    paths = [param_path(p) for p in raw_paths]
    return paths, list(leaves), treedef

tree_leaves(tree)

Extract all leaf tensors from a PyTree.

Parameters:

Name Type Description Default
tree Any

Nested structure (dict/list/tuple/…) containing tensors as leaves.

 required

Returns:

Type Description
list[Tensor]

List of all tensor leaves in the tree (non-tensor leaves are ignored).

Raises:

Type Description
InputTypeError

If tree contains a DP wrapper.
Example

tree = {'a': torch.tensor([1, 2]), 'b': {'c': torch.tensor([3])}} leaves = tree_leaves(tree) len(leaves) 2

Source code in packages/opake-engine/src/opake/api/engine/pytree.py 
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
def tree_leaves(tree: Any) -> list[torch.Tensor]:
    """Extract all leaf tensors from a PyTree.

    Args:
        tree: Nested structure (dict/list/tuple/…) containing tensors as leaves.

    Returns:
        List of all tensor leaves in the tree (non-tensor leaves are ignored).

    Raises:
        InputTypeError: If ``tree`` contains a DP wrapper.

    Example:
        >>> tree = {'a': torch.tensor([1, 2]), 'b': {'c': torch.tensor([3])}}
        >>> leaves = tree_leaves(tree)
        >>> len(leaves)
        2
    """
    flat, _ = _ot.tree_flatten(tree)
    tensor_leaves = [leaf for leaf in flat if isinstance(leaf, torch.Tensor)]
    if len(tensor_leaves) == len(flat):
        return tensor_leaves

    # Local import avoids the types -> pytree cycle.
    from opake.api.engine.types import ClippedPytree

    if any(isinstance(leaf, ClippedPytree) for leaf in flat):
        raise InputTypeError(_RAW_TENSOR_PYTREE_ERROR)
    return tensor_leaves

tree_map(fn, *trees)

Apply function to all leaves of one or more PyTrees.

Parameters:

Name Type Description Default
fn Callable[..., Any]

Function to apply to each leaf (or set of leaves from multiple trees).

 required
*trees Any

One or more PyTrees with matching structure.

 ()

Returns:

Type Description
Any

PyTree with same structure as inputs, with fn applied to leaves.
Example

tree = {'a': torch.tensor([1.0, 2.0]), 'b': torch.tensor([3.0])} doubled = tree_map(lambda x: x * 2, tree) doubled['a'] tensor([2., 4.])

Source code in packages/opake-engine/src/opake/api/engine/pytree.py 
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
def tree_map(fn: Callable[..., Any], *trees: Any) -> Any:
    """Apply function to all leaves of one or more PyTrees.

    Args:
        fn: Function to apply to each leaf (or set of leaves from multiple trees).
        *trees: One or more PyTrees with matching structure.

    Returns:
        PyTree with same structure as inputs, with `fn` applied to leaves.

    Example:
        >>> tree = {'a': torch.tensor([1.0, 2.0]), 'b': torch.tensor([3.0])}
        >>> doubled = tree_map(lambda x: x * 2, tree)
        >>> doubled['a']
        tensor([2., 4.])
    """
    return _ot.tree_map(fn, *trees)

tree_map_with_path(fn, tree)

Apply function to leaves with their :data:ParamPath in the tree.

Paths and leaf order match :func:tree_flatten_with_paths (optree traversal, including sorted dict keys), so results align with :class:~opake.types.PerGroup keys.

Parameters:

Name Type Description Default
fn Callable[[ParamPath, Any], Any]

Function that takes (path, leaf) where path is a :data:ParamPath.

 required
tree Any

PyTree to traverse

 required

Returns:

Type Description
Any

PyTree with same structure, with fn applied to (path, leaf)
Example

tree = {'layer1': {'weight': torch.ones(2)}, 'layer2': {'bias': torch.zeros(3)}} def print_shapes(path, leaf): ... print(f"{path}: {leaf.shape}") ... return leaf tree_map_with_path(print_shapes, tree) ('layer1', 'weight'): torch.Size([2]) ('layer2', 'bias'): torch.Size([3])

Source code in packages/opake-engine/src/opake/api/engine/pytree.py 
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
def tree_map_with_path(
    fn: Callable[[ParamPath, Any], Any],
    tree: Any,
) -> Any:
    """Apply function to leaves with their :data:`ParamPath` in the tree.

    Paths and leaf order match :func:`tree_flatten_with_paths` (optree
    traversal, including sorted dict keys), so results align with
    :class:`~opake.types.PerGroup` keys.

    Args:
        fn: Function that takes ``(path, leaf)`` where ``path`` is a
            :data:`ParamPath`.
        tree: PyTree to traverse

    Returns:
        PyTree with same structure, with fn applied to (path, leaf)

    Example:
        >>> tree = {'layer1': {'weight': torch.ones(2)}, 'layer2': {'bias': torch.zeros(3)}}
        >>> def print_shapes(path, leaf):
        ...     print(f"{path}: {leaf.shape}")
        ...     return leaf
        >>> tree_map_with_path(print_shapes, tree)
        ('layer1', 'weight'): torch.Size([2])
        ('layer2', 'bias'): torch.Size([3])
    """
    paths, leaves, treedef = tree_flatten_with_paths(tree)
    out = [fn(path, leaf) for path, leaf in zip(paths, leaves, strict=True)]
    return _ot.tree_unflatten(treedef, out)

tree_structure(tree)

Return the optree structure of tree (no leaves).

Thin wrapper around :func:optree.tree_structure. Useful for asserting that independently gathered payloads share a layout before unflattening.

Source code in packages/opake-engine/src/opake/api/engine/pytree.py 
108
109
110
111
112
113
114
def tree_structure(tree: Any) -> Any:
    """Return the optree structure of ``tree`` (no leaves).

    Thin wrapper around :func:`optree.tree_structure`.  Useful for asserting
    that independently gathered payloads share a layout before unflattening.
    """
    return _ot.tree_structure(tree)

tree_unflatten(treedef, leaves)

Rebuild a PyTree from treedef and leaves.

Thin wrapper around :func:optree.tree_unflatten.

Source code in packages/opake-engine/src/opake/api/engine/pytree.py 
100
101
102
103
104
105
def tree_unflatten(treedef: Any, leaves: list[Any]) -> Any:
    """Rebuild a PyTree from ``treedef`` and ``leaves``.

    Thin wrapper around :func:`optree.tree_unflatten`.
    """
    return _ot.tree_unflatten(treedef, leaves)