"""TarFlow: transformer autoregressive normalizing flow.
Adapted from apple/ml-tarflow (TarFlow) and apple/ml-starflow (STARFlow);
see models/tarflow/LICENSE.apple and LICENSE.starflow.
Self-contained ``(B, T, F)`` density model (absorbs the former
``TransformerFlow`` container and the ``make_tarflow`` factory). Head sizing
follows the Flux1 convention: specify ``head_dim`` and ``num_heads``; total
width ``channels = head_dim * num_heads`` is derived.
"""
from dataclasses import dataclass
import jax
import jax.numpy as jnp
from flax import nnx
from jax import Array
from jax.typing import DTypeLike
from gensbi.models.core.stats import fit_stat
from gensbi.models.core.tokenizers import VectorTokenizer, ImageTokenizer
from gensbi.models.tarflow.blocks import MetaBlock
from gensbi.models.tarflow.conditioners import (
AdditiveBiasConditioner, VectorConditioner, ImageConditioner,
)
from gensbi.models.tarflow.pe import VisionRotaryEmbedding
from gensbi.normalizing_flows.bijections.base import Mask
[docs]
_LOG2PI = jnp.log(2.0 * jnp.pi)
@dataclass
[docs]
class TarFlowParams:
"""Architecture parameters for :class:`TarFlow`.
``modeled`` selects the tokenizer (``"vector"`` or ``"image"``); ``cond``
selects the conditioner (``"bias"``, ``"vector"``, or
``"image"``). Head sizing follows the Flux1 convention: specify
``head_dim`` and ``num_heads``; total width
``channels = head_dim * num_heads`` is derived in ``__post_init__``.
Parameters
----------
rngs : nnx.Rngs
Flax RNG container passed to all sub-modules during construction.
dim : int or None, optional
Feature dimension of each input vector. Required when
``modeled="vector"``. Default is ``None``.
cond_dim : int, optional
Dimensionality of the conditioning vector. Set to ``0`` for an
unconditional model. Default is ``0``.
modeled : str, optional
Tokenizer type: ``"vector"`` (1-D data) or ``"image"`` (spatial data).
Default is ``"vector"``.
img_size : int or None, optional
Spatial size (height = width) of the modeled image. Required when
``modeled="image"``. Default is ``None``.
patch_size : int or None, optional
Patch size for the image tokenizer. Required when
``modeled="image"``. Default is ``None``.
img_channels : int, optional
Number of channels in the modeled image. Default is ``1``.
cond : str, optional
Conditioning strategy: ``"bias"`` (per-token additive bias via
:class:`~gensbi.models.tarflow.conditioners.AdditiveBiasConditioner`),
``"vector"`` (one condition token per modeled coordinate via
:class:`~gensbi.models.tarflow.conditioners.VectorConditioner`),
or ``"image"`` (prefix tokens from an image via
:class:`~gensbi.models.tarflow.conditioners.ImageConditioner`).
Default is ``"bias"``.
cond_img_size : int or None, optional
Spatial size of the conditioning image. Required when
``cond="image"``. Default is ``None``.
cond_patch_size : int or None, optional
Patch size for the image conditioning tokenizer. Required when
``cond="image"``. Default is ``None``.
cond_channels : int, optional
Number of channels in the conditioning image. Default is ``1``.
head_dim : int, optional
Dimension per attention head. Default is ``16``.
num_heads : int, optional
Number of attention heads per block. Default is ``4``.
num_blocks : int, optional
Number of :class:`~gensbi.models.tarflow.blocks.MetaBlock` layers.
Default is ``8``.
layers_per_block : int, optional
Number of :class:`~gensbi.models.tarflow.blocks.AttentionBlock`
layers inside each :class:`~gensbi.models.tarflow.blocks.MetaBlock`.
Default is ``2``.
block_size : int, optional
Token grouping factor for the vector tokenizer. Default is ``1``.
permutation : str, optional
Token permutation strategy per block: ``"flip"`` (alternate
forward/reverse order) or ``"random"`` (independently sampled per
block). Default is ``"flip"``.
standardize : bool, optional
If ``True`` (default), apply mean/std standardization to inputs and
outputs. Enables :meth:`TarFlow.set_standardization`.
zero_init : bool, optional
If ``True`` (default), initialize ``proj_out`` weights to zero so
each :class:`~gensbi.models.tarflow.blocks.MetaBlock` starts as the
identity map.
use_softplus : bool, optional
If ``True`` (default), use softplus for the affine scale (numerically
stable, bounded tail). If ``False``, use ``exp`` (legacy behavior).
soft_clip : float, optional
Soft-clip magnitude applied via ``tanh`` to raw network outputs before
splitting into ``(a, b)``. Default is ``4.0``.
use_rope : bool, optional
If ``True``, replace the learned per-token ``pos_embed`` for the
modeled image tokens with 2D rotary position embeddings
(:class:`~gensbi.models.tarflow.pe.VisionRotaryEmbedding`). Prefix
(condition) tokens keep their learned embeddings and sit at the
identity rotation (zero angles). Requires ``modeled="image"`` and
``head_dim`` divisible by ``4``. ``head_dim >= 32`` is recommended
for image data (more rotary frequencies per axis) but not enforced.
Default is ``False``.
rope_theta : int, optional
Frequency base for the rotary embedding. Only used when
``use_rope=True``. Default is ``10000``.
param_dtype : DTypeLike, optional
Dtype for all stored (master) kernel/bias/embedding parameters
across the tokenizer, conditioner, and transformer blocks. Default
is ``float32``.
dtype : DTypeLike, optional
Compute dtype knob threaded through the conditioners and
:class:`~gensbi.models.tarflow.blocks.MetaBlock`/
:class:`~gensbi.models.tarflow.blocks.AttentionBlock` layers.
Default is ``float32``, matching ``param_dtype``, so with default
arguments this is a bit-identical no-op cast. Hard-fp32 regardless
of this knob: the softplus/soft_clip affine-scale path in
``MetaBlock._affine``, log-det accumulation, the ``mean``/``std``
standardization buffers, and the KV-cache buffers used during
sampling.
"""
[docs]
modeled: str = "vector"
[docs]
img_size: int | None = None
[docs]
patch_size: int | None = None
[docs]
cond_img_size: int | None = None
[docs]
cond_patch_size: int | None = None
[docs]
layers_per_block: int = 2
[docs]
permutation: str = "flip"
[docs]
standardize: bool = True
[docs]
use_softplus: bool = True
[docs]
rope_theta: int = 10000
[docs]
param_dtype: DTypeLike = jnp.float32
[docs]
dtype: DTypeLike = jnp.float32
[docs]
def __post_init__(self):
if self.modeled not in ("vector", "image"):
raise ValueError(f"unknown modeled {self.modeled!r}")
if self.modeled == "vector" and self.dim is None:
raise ValueError("modeled='vector' requires dim")
if self.modeled == "image" and (self.img_size is None or self.patch_size is None):
raise ValueError("modeled='image' requires img_size and patch_size")
if self.cond not in ("bias", "vector", "image"):
raise ValueError(f"unknown cond {self.cond!r}")
if self.cond == "image" and (self.cond_img_size is None or self.cond_patch_size is None):
raise ValueError("cond='image' requires cond_img_size and cond_patch_size")
if self.permutation not in ("flip", "random"):
raise ValueError(f"unknown permutation {self.permutation!r}")
if self.use_rope:
if self.modeled != "image":
raise ValueError("use_rope=True requires modeled='image'")
if self.head_dim % 4 != 0:
raise ValueError(
f"use_rope requires head_dim divisible by 4 (two position "
f"axes x adjacent-pair rotation), got head_dim={self.head_dim}")
self.channels = self.head_dim * self.num_heads
[docs]
class TarFlow(nnx.Module):
"""Transformer autoregressive normalizing flow density model.
Stacks :class:`~gensbi.models.tarflow.blocks.MetaBlock` bijections with
alternating token permutations on top of a tokenizer and an isotropic
Gaussian base distribution. Supports both vector and image data, with
optional input standardization.
Parameters
----------
params : TarFlowParams
Architecture and initialization parameters.
"""
def __init__(self, params: TarFlowParams):
rngs = params.rngs
channels = params.channels
if params.modeled == "vector":
tokenizer = VectorTokenizer(params.dim, params.block_size,
params.vec_channels)
else:
tokenizer = ImageTokenizer(params.img_size, params.img_size,
params.img_channels, params.patch_size)
T, F = tokenizer.T, tokenizer.F
if params.use_rope:
# tokenizer.grid[0] assumes a square grid; always true today since
# ImageTokenizer is built with img_size x img_size (see above).
rope = VisionRotaryEmbedding(dim=params.head_dim // 2,
pt_seq_len=tokenizer.grid[0],
theta=params.rope_theta)
grid = tokenizer.grid
else:
rope, grid = None, None
def make_cond():
if params.cond == "bias":
return AdditiveBiasConditioner(params.cond_dim, channels, rngs=rngs,
cond_channels=params.cond_channels,
param_dtype=params.param_dtype,
dtype=params.dtype)
if params.cond == "vector":
return VectorConditioner(params.cond_dim, params.cond_channels,
channels, rngs=rngs,
param_dtype=params.param_dtype,
dtype=params.dtype)
m = (params.cond_img_size // params.cond_patch_size) ** 2
return ImageConditioner(params.cond_channels,
params.cond_patch_size, channels, m,
rngs=rngs, param_dtype=params.param_dtype,
dtype=params.dtype)
blocks = []
for i in range(params.num_blocks):
if params.permutation == "flip":
perm = jnp.arange(T) if i % 2 == 0 else jnp.arange(T)[::-1]
else:
perm = jax.random.permutation(rngs.params(), T)
blocks.append(MetaBlock(
F=F, channels=channels, T=T, perm=perm,
conditioner=make_cond(), num_layers=params.layers_per_block,
num_heads=params.num_heads, expansion=4, rngs=rngs,
zero_init=params.zero_init, use_softplus=params.use_softplus,
soft_clip=params.soft_clip, rope=rope, grid=grid,
param_dtype=params.param_dtype, dtype=params.dtype))
[docs]
self.blocks = nnx.List(blocks)
[docs]
self.tokenizer = tokenizer
[docs]
self.cond_dim = params.cond_dim
[docs]
self.example_shape = tokenizer.example_shape
[docs]
self._standardize = params.standardize
[docs]
self.mean = Mask(jnp.zeros(self.example_shape))
[docs]
self.std = Mask(jnp.ones(self.example_shape))
[docs]
def _base_log_prob(self, z: Array) -> Array:
return -0.5 * jnp.sum(z ** 2, axis=(1, 2)) - 0.5 * self.T * self.F * _LOG2PI
[docs]
def _ensure_batched(self, x: Array) -> Array:
x = jnp.asarray(x)
if x.ndim == len(self.example_shape):
x = x[None]
return x
[docs]
def log_prob(self, x: Array, cond: Array | None = None) -> Array:
"""Compute the log-probability of data under the model.
Applies standardization, tokenizes the input, then runs each
:class:`~gensbi.models.tarflow.blocks.MetaBlock`'s
:meth:`~gensbi.models.tarflow.blocks.MetaBlock.inverse` transform
(data→noise direction), accumulating the log-absolute-determinant
terms, and finally evaluates the base Gaussian log-probability.
Parameters
----------
x : Array
Data samples of shape ``(B, *example_shape)`` or a single
unbatched sample that will be promoted to a batch of one.
cond : Array or None, optional
Conditioning batch of shape ``(B, cond_dim)`` for
``cond_channels == 1``, or ``(B, cond_dim, C_cond)`` for
``cond_channels > 1`` (flattened internally by the conditioner).
Pass ``None`` for an unconditional model.
Returns
-------
Array
Log-probabilities of shape ``(B,)``.
"""
x = self._ensure_batched(x)
u = (x - self.mean[...]) / self.std[...]
logdet = -jnp.sum(jnp.log(self.std[...]))
z = self.tokenizer.tokenize(u)
total = jnp.broadcast_to(logdet, (x.shape[0],))
for blk in self.blocks:
z, ld = blk.inverse(z, cond)
total = total + ld
return self._base_log_prob(z) + total
[docs]
def sample(self, key, cond: Array | None = None, nsamples: int | None = None):
"""Draw samples from the model.
Samples noise from ``N(0, I)``, then applies each
:class:`~gensbi.models.tarflow.blocks.MetaBlock`'s
:meth:`~gensbi.models.tarflow.blocks.MetaBlock.forward` transform
(noise→data direction) in reverse block order, detokenizes the
result, and applies the inverse standardization.
Parameters
----------
key : jax.random.PRNGKey
Random key for noise sampling.
cond : Array or None, optional
Conditioning batch of shape ``(nsamples, cond_dim)`` for
``cond_channels == 1``, or ``(nsamples, cond_dim, C_cond)`` for
``cond_channels > 1`` (flattened internally by the conditioner).
If provided, ``nsamples`` is inferred from ``cond.shape[0]``.
nsamples : int or None, optional
Number of samples to draw. Required when ``cond`` is ``None``.
Returns
-------
Array
Samples of shape ``(B, *example_shape)``.
"""
if cond is not None:
nsamples = cond.shape[0]
z = jax.random.normal(key, (nsamples, self.T, self.F))
x = z
for blk in reversed(self.blocks):
x, _ = blk.forward(x, cond)
x = self.tokenizer.detokenize(x)
return x * self.std[...] + self.mean[...]
[docs]
def set_standardization(self, mean, std) -> None:
"""Set the mean and standard deviation for input standardization.
Accepts shapes ``(dim,)`` (broadcast to ``(dim, 1)``), ``(dim, 1)``,
``(C,)`` (per-channel broadcast), or a scalar broadcastable to
``example_shape``.
Parameters
----------
mean : Array
Mean broadcastable to ``example_shape``.
std : Array
Standard deviation broadcastable to ``example_shape``.
Returns
-------
None
Raises
------
ValueError
If the model was built with ``standardize=False``.
"""
if not self._standardize:
raise ValueError("TarFlow built with standardize=False")
self.mean[...] = fit_stat(mean, self.example_shape, dtype=self.mean[...].dtype)
self.std[...] = fit_stat(std, self.example_shape, dtype=self.std[...].dtype)
