class

CatTransform

extendsTransform

CatTransform(transforms: list[Transform], dim: int = 0, lengths: list[int] | None = None)

source

Apply different transforms to contiguous partitions along an axis.

Splits the input along dim into contiguous chunks of size lengths[i] (or equal partitions if lengths is None), applies the $i$ -th transform to the $i$ -th chunk, and concatenates the results back along the same axis. Differs from StackTransform in that each partition can have a different length, not just a single index.

Parameters

transformslist[Transform]

One transform per partition.

dimint= 0

Concatenation dimension. Default 0.

lengthslist[int]= None

Length of each partition. If None the axis size must be divisible by len(transforms) and equal partitions are used.

Raises

ValueError

If transforms is empty, or if the axis size is not divisible by len(transforms) when lengths is None.

Notes

Forward (with $\mathbf{x}_i$ the $i$ -th partition):

\mathbf{y} = \mathrm{cat}\bigl([T_i(\mathbf{x}_i)]_{i=1}^{n}, \;\mathrm{dim}\bigr)

Inverse: split, invert, re-concatenate.

Log Jacobian determinant: per-partition Jacobians concatenated back along dim:

\log|\det J| = \mathrm{cat}\bigl( [\log|\det J_{T_i}|(\mathbf{x}_i, \mathbf{y}_i)]_{i=1}^{n}, \;\mathrm{dim}\bigr)

Examples

>>> import lucid
>>> from lucid.distributions.transforms import ExpTransform, AffineTransform, CatTransform
>>> T = CatTransform([ExpTransform(), AffineTransform(0.0, 2.0)],
...                  dim=0, lengths=[2, 3])
>>> T(lucid.tensor([0.0, 0.0, 1.0, 2.0, 3.0])).shape
(5,)

Methods (2)

dunder

init

→None

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

source

Store the per-partition transforms, concat axis, and partition lengths.

Raises

ValueError

If transforms is empty.

log_abs_det_jacobian

→Tensor

log_abs_det_jacobian(x: Tensor, y: Tensor)

source

Per-partition Jacobians concatenated back along dim.