bijx.BetaStretch

class bijx.BetaStretch

Bases: ScalarBijection

Beta-inspired stretching on unit interval.

Applies a nonlinear stretching transformation on the unit interval inspired by the Beta distribution. This bijection can compress or expand different regions of [0,1] depending on the parameter value.

Type: \([0, 1] \to [0, 1]\)

Transform: \(f(x) = \frac{x^a}{x^a + (1-x)^a}\) where \(a > 0\)

Parameters:

• a (Union[Variable, Array, ndarray, Sequence[Union[int, Any]]]) – Shape parameter controlling the stretching behavior.

• transform_a (Callable | None) – Function to ensure positive parameter (default: softplus).

• rngs – Random number generators for parameter initialization.

Note

• When a = 1, this is the identity transformation

• When a < 1, the transformation is S-shaped (sigmoidal)

• When a > 1, the transformation is reverse S-shaped

• The parameter is initialized to give a = 1 after transformation

Example

>>> bijection = BetaStretch(rngs=rngs)
>>> x = jnp.array([0.25, 0.5, 0.75])
>>> y, log_det = bijection.forward(x, jnp.zeros(3))

__init__(a=(), transform_a=<PjitFunction of <function softplus>>, *, rngs=None)

Parameters:

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

• transform_a (Callable | None)

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\).