Leafwise

class pmrf.stats.Leafwise(bijectors: PyTree[distreqx.bijectors._bijector.AbstractBijector])

Bases: AbstractFwdLogDetJacBijector, AbstractInvLogDetJacBijector

Applies a pytree of bijectors to a pytree of inputs.

This behaves analogously to TensorFlow Probability’s JointMap. It allows applying independent bijectors to a structured input (e.g., a tuple or dict of arrays) and aggregates the log-determinants across the structure.

Initializes a TreeMap bijector.

forward(x: PyTree) → PyTree

Computes y = f(x).

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

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

inverse(y: PyTree) → PyTree

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

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

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

same_as(other: AbstractBijector) → bool

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

bijectors: PyTree[distreqx.bijectors._bijector.AbstractBijector]