gensbi.normalizing_flows.bijections.base#

Bijection abstract base and the non-trainable Mask variable type.

Classes#

Bijection

Invertible map between a noise variable and a data variable.

Mask

Fixed (non-trainable) buffer for bijection metadata.

Module Contents#

class gensbi.normalizing_flows.bijections.base.Bijection[source]#

Bases: flax.nnx.Module

Invertible map between a noise variable and a data variable.

Subclasses implement a differentiable bijection with a direction convention fixed across the library: forward() maps noise to data (the sampling direction) and inverse() maps data to noise (the density-evaluation direction). Both directions also return the log absolute determinant of the Jacobian of the transform they apply.

abstractmethod forward(u, cond=None)[source]#

Map noise to data (the sampling direction).

Parameters:

  • u (Array) – Noise-space input.

  • cond (Array or None, optional) – Conditioning input, or None for an unconditional map.

Returns:

  • x (Array) – Data-space output.

  • logabsdet (Array) – Log absolute determinant of the Jacobian of the forward map.

Raises:

NotImplementedError – This is an abstract method; subclasses must override it.

Return type:

tuple[jax.Array, jax.Array]

abstractmethod inverse(x, cond=None)[source]#

Map data to noise (the density-evaluation direction).

Parameters:

  • x (Array) – Data-space input.

  • cond (Array or None, optional) – Conditioning input, or None for an unconditional map.

Returns:

  • u (Array) – Noise-space output.

  • logabsdet (Array) – Log absolute determinant of the Jacobian of the inverse map.

Raises:

NotImplementedError – This is an abstract method; subclasses must override it.

Return type:

tuple[jax.Array, jax.Array]

class gensbi.normalizing_flows.bijections.base.Mask(value, *, hijax=None, ref=None, eager_sharding=None, **metadata)[source]#

Bases: flax.nnx.Variable

Fixed (non-trainable) buffer for bijection metadata.

A non-Param Variable subclass so that optimizers and EMA utilities skip it while checkpointing still saves and restores it. Typical use: autoregressive masks in MaskedAutoregressive.

Parameters:

  • value (A | VariableMetadata[A])

  • hijax (bool | None)

  • ref (bool | None)

  • eager_sharding (bool | None)

  • metadata (Any)