AbstractBijector

class pmrf.stats.bijectors.AbstractBijector

Bases: AbstractStrictModule

Differentiable bijection that knows to compute its Jacobian determinant.

A bijector implements a differentiable and bijective transformation f, whose inverse is also differentiable (f is called a “diffeomorphism”). A bijector can be used to transform a continuous random variable X to a continuous random variable Y = f(X) in the context of TransformedDistribution.

Typically, a bijector subclass will implement the following methods:

• forward_and_log_det(x) (required)

• inverse_and_log_det(y) (optional)

The remaining methods are defined in terms of the above by default.

Subclass requirements:

• Subclasses must ensure that f is differentiable and bijective, and that their methods correctly implement f^{-1}, J(f) and J(f^{-1}). Distreqx will assume these properties hold, and will make no attempt to verify them.

abstractmethod forward(x: PyTree) → PyTree

Computes \(y = f(x)\).

abstractmethod forward_and_log_det(x: PyTree) → tuple[PyTree, PyTree]

Computes \(y = f(x)\) and \(\log|\det J(f)(x)|\).

abstractmethod forward_log_det_jacobian(x: PyTree) → PyTree

Computes \(\log|\det J(f)(x)|\).

abstractmethod inverse(y: PyTree) → PyTree

Computes \(x = f^{-1}(y)\).

abstractmethod inverse_and_log_det(y: Array) → tuple[PyTree, PyTree]

Computes \(x = f^{-1}(y)\) and \(\log|\det J(f^{-1})(y)|\).

abstractmethod inverse_log_det_jacobian(y: PyTree) → PyTree

Computes \(\log|\det J(f^{-1})(y)|\).

abstractmethod same_as(other) → bool

Returns True if this bijector is guaranteed to be the same as other.

property is_constant_jacobian: bool

Whether the Jacobian is promised to be constant.

property is_constant_log_det: bool

Whether the Jacobian determinant is promised to be constant.

property name: str

Name of the bijector.