bijx.AffineLinear

class bijx.AffineLinear

Bases: ScalarBijection

Learnable affine transformation.

Applies a learnable affine transformation combining scaling and shifting. This is one of the most fundamental bijections, providing location-scale transformations commonly used in normalizing flows.

Type: \([-\infty, \infty] \to [-\infty, \infty]\)

Transform: \(ax + b\) where \(a\) (scale) and \(b\) (shift) are learnable parameters

Parameters:

• scale (Union[Variable, Array, ndarray, Sequence[Union[int, Any]]]) – Scale parameter specification, transformed to ensure appropriate scaling.

• shift (Union[Variable, Array, ndarray, Sequence[Union[int, Any]]]) – Shift parameter specification, no transformation by default.

• transform_scale (Callable | None) – Function to transform scale (default: exp for positivity).

• transform_shift (Callable | None) – Function to transform shift (default: None/identity).

• rngs (Rngs) – Random number generators for parameter initialization.

Example

>>> bijection = AffineLinear(rngs=rngs)
>>> x = jnp.array([-1.0, 0.0, 1.0])
>>> y, log_det = bijection.forward(x, jnp.zeros(3))

__init__(scale=(), shift=(), transform_scale=<PjitFunction of <function exp>>, transform_shift=None, *, rngs=None)

Parameters:

• scale (Variable | Array | ndarray | Sequence[int | Any])

• shift (Variable | Array | ndarray | Sequence[int | Any])

• transform_scale (Callable | None)

• transform_shift (Callable | None)

• rngs (Rngs)

Methods

forward(x, log_density, **kwargs)	Apply forward transformation with log-density update.
fwd(x, **kwargs)	Apply forward transformation.
invert()	Create an inverted version of this bijection.
log_jac(x, y, **kwargs)	Compute log absolute determinant of the Jacobian.
rev(y, **kwargs)	Apply reverse (inverse) transformation.
reverse(y, log_density, **kwargs)	Apply reverse transformation with log-density update.

log_jac(x, y, **kwargs)

Compute log absolute determinant of the Jacobian.

Parameters:

• x – Input values where Jacobian is computed.

• y – Output values corresponding to x (i.e., y = fwd(x)).

• **kwargs – Additional transformation-specific arguments.

Returns:

Log absolute Jacobian determinant \(\log \abs{f'(x)}\) with same shape as x.

fwd(x, **kwargs)

Apply forward transformation.

Parameters:

• x – Input values to transform.

• **kwargs – Additional transformation-specific arguments.

Returns:

Transformed values \(y = f(x)\) with same shape as \(x\).

rev(y, **kwargs)

Apply reverse (inverse) transformation.

Parameters:

• y – Output values to inverse-transform.

• **kwargs – Additional transformation-specific arguments.

Returns:

Inverse-transformed values \(x = f^{-1}(y)\) with same shape as \(y\).