class

StackTransform

extendsTransform

StackTransform(transforms: list[Transform], dim: int = 0)

source

Apply a list of transforms to indexed slices along a stack axis.

The $k$ -th transform is applied to the $k$ -th slice obtained by Tensor.unbind along dim, and the results are stacked back along the same axis. This is the heterogeneous counterpart of IndependentTransform: different bijections can be applied to different coordinates of an event vector.

Parameters

transformslist[Transform]

One transform per slice. Must be non-empty and the same length as x.shape[dim].

dimint= 0

Dimension along which to slice. Default 0.

Raises

ValueError

If transforms is empty, or if the size of dim does not match len(transforms) at call time.

Notes

Forward (with $\mathbf{x}_k = \mathbf{x}[\ldots, k, \ldots]$ ):

\mathbf{y}_k = T_k(\mathbf{x}_k)

Inverse:

\mathbf{x}_k = T_k^{-1}(\mathbf{y}_k)

Log Jacobian determinant (per-slice values stacked back along dim):

\log|\det J|_k = \log|\det J_{T_k}|(\mathbf{x}_k, \mathbf{y}_k)

Useful for "block" flows where each coordinate of the latent vector has its own bijection — e.g., applying ExpTransform to positive components and TanhTransform to bounded ones.

Examples

>>> import lucid
>>> from lucid.distributions.transforms import ExpTransform, TanhTransform, StackTransform
>>> T = StackTransform([ExpTransform(), TanhTransform()], dim=-1)
>>> T(lucid.tensor([0.0, 0.0]))  # (exp(0), tanh(0)) = (1, 0)
Tensor([1.0, 0.0])

Methods (2)

dunder

init

→None

__init__(transforms: list[Transform], dim: int = 0)

source

Store the per-slice transforms and the stacking axis dim.

Raises

ValueError

If transforms is empty.

log_abs_det_jacobian

→Tensor

log_abs_det_jacobian(x: Tensor, y: Tensor)

source

Per-slice Jacobians stacked back along dim.