gensbi.normalizing_flows.bijections.permutation#

Permutation bijection (dimension reordering between flow layers).

Classes#

Reorder dimensions; conditioning is ignored; log-det is always 0.

class gensbi.normalizing_flows.bijections.permutation.Permutation(perm)[source]#

Reorder dimensions; conditioning is ignored; log-det is always 0.

Both the permutation index array and its inverse are stored as Mask buffers so that optimizers and EMA utilities skip them.

Parameters:: perm (Array) – Integer index array of shape (dim,) that defines the reordering. perm[i] is the source index for output position i.

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

Map noise to data by applying the inverse permutation.

Parameters:

u (Array) – Noise-space input of shape (dim,).
cond (Array or None, optional) – Ignored; present for interface compatibility.

Returns:

x (Array) – Reordered data-space output.
logabsdet (Array) – Zero scalar (permutations have unit Jacobian determinant).

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

Map data to noise by applying perm to reorder dimensions.

Parameters:

x (Array) – Data-space input of shape (dim,).
cond (Array or None, optional) – Ignored; present for interface compatibility.

Returns:

u (Array) – Reordered noise-space output.
logabsdet (Array) – Zero scalar (permutations have unit Jacobian determinant).

classmethod random(dim, rngs)[source]#

Construct a uniformly random permutation.

Parameters:

dim (int) – Number of dimensions.
rngs (nnx.Rngs) – Flax RNG container used to draw the random permutation index.

Returns:

A Permutation with a randomly shuffled index array.

Return type:

Permutation

classmethod reverse(dim)[source]#

Construct a permutation that reverses dimension order.