import abc
import math
from collections.abc import Callable, Sequence
from typing import ClassVar
import jax
import jax.numpy as jnp
from flax import linen as nn
from flax.linen import initializers
from cellflow._types import Layers_t
__all__ = [
"SelfAttention",
"SeedAttentionPooling",
"TokenAttentionPooling",
"MLPBlock",
"FilmBlock",
"ResNetBlock",
"SelfAttentionBlock",
"sinusoidal_time_encoder",
"LAYER_REGISTRY",
"register_layer",
]
def sinusoidal_time_encoder(t: jnp.ndarray, time_freqs: int = 1024, time_max_period: int | None = 10000) -> jnp.ndarray:
"""
Create sinusoidal timestep embeddings.
Parameters
----------
t : jnp.ndarray
A 1-D or 2-D array of timesteps. May be fractional.
num_freqs : int
The dimension of the embedding.
max_period : int | None
Controls the minimum frequency of the embeddings.
If :obj:`None`, the frequencies are evenly spaced in the range [0, 2 * pi],
reimplementing :func:`ott.neural.networks.layers.cyclical_time_encoder`
for backward compatibility.
Returns
-------
jnp.ndarray
Sinusoidal embedding.
"""
if time_max_period is None:
freq = 2 * jnp.arange(time_freqs) * jnp.pi
t = freq * t
return jnp.concatenate([jnp.cos(t), jnp.sin(t)], axis=-1)
t = t * time_max_period
half = time_freqs // 2
freqs = jnp.exp(-math.log(time_max_period) * jnp.arange(start=0, stop=half, dtype=jnp.float32) / half)
args = t * freqs
embedding = jnp.concatenate([jnp.cos(args), jnp.sin(args)], axis=-1)
return embedding
[docs]
class BaseModule(abc.ABC, nn.Module):
"""Base module for condition encoder and its components."""
#: Whether this block's ``__call__`` accepts an attention mask as its second
#: positional argument. Declared as a :class:`~typing.ClassVar` so it stays a
#: pure class-level capability flag and is *not* turned into a ``flax`` module
#: field (which would alter the constructor signature). :func:`_apply_modules`
#: dispatches on this flag instead of hard-coding ``isinstance`` checks.
takes_attention_mask: ClassVar[bool] = False
[docs]
@abc.abstractmethod
def __call__(self, x: jnp.ndarray, training: bool = True) -> jnp.ndarray:
"""Forward pass."""
pass
#: Registry mapping a ``layer_type`` string to the block class that implements it.
#: Populated via the :func:`register_layer` decorator and consulted by
#: :func:`_get_layers` to instantiate blocks from config dicts. New block types can
#: be added by decorating them with ``@register_layer(...)`` without touching the
#: dispatch code.
LAYER_REGISTRY: dict[str, type[BaseModule]] = {}
def register_layer(name: str) -> Callable[[type[BaseModule]], type[BaseModule]]:
"""Register a network block class under ``name`` in :data:`LAYER_REGISTRY`.
The decorated class can then be referenced by ``name`` through the
``"layer_type"`` key of a layer config and is instantiated by
:func:`_get_layers`.
Parameters
----------
name
Key under which the decorated class is registered.
Returns
-------
A class decorator that registers the class and returns it unchanged.
"""
def decorator(cls: type[BaseModule]) -> type[BaseModule]:
LAYER_REGISTRY[name] = cls
return cls
return decorator
[docs]
@register_layer("mlp")
class MLPBlock(BaseModule):
"""
MLP block.
Parameters
----------
dims
Dimensions of the MLP layers.
dropout_rate
Dropout rate.
act_last_layer
Whether to apply the activation function to the last layer.
act_fn
Activation function.
"""
dims: Sequence[int] = (1024, 1024, 1024)
dropout_rate: float = 0.0
act_last_layer: bool = True
act_fn: Callable[[jnp.ndarray], jnp.ndarray] = nn.silu
[docs]
@nn.compact
def __call__(self, x: jnp.ndarray, training: bool = True) -> jnp.ndarray:
"""
Forward pass.
Parameters
----------
x
Input tensor of shape ``(batch_size, input_dim)``.
training
Whether the model is in training mode.
Returns
-------
Output tensor of shape ``(batch_size, output_dim)``.
"""
if len(self.dims) == 0:
return x
z = x
for i in range(len(self.dims) - 1):
z = self.act_fn(nn.Dense(self.dims[i])(z))
z = nn.Dropout(self.dropout_rate)(z, deterministic=not training)
z = nn.Dense(self.dims[-1])(z)
z = self.act_fn(z) if self.act_last_layer else z
z = nn.Dropout(self.dropout_rate)(z, deterministic=not training)
return z
[docs]
class FilmBlock(BaseModule):
"""Feature-wise Linear Modulation (FiLM) layer.
Applies a learned affine transformation (scale and shift) to the input,
conditioned on an external embedding.
Parameters
----------
input_dim
Dimensionality of the input features.
cond_dim
Dimensionality of the conditioning features.
act_fn
Activation function to apply after modulation. If :obj:`None`, no activation is applied.
"""
input_dim: int
cond_dim: int
act_fn: Callable[[jnp.ndarray], jnp.ndarray] = lambda x: x
[docs]
def setup(self) -> None:
self.film_generator = nn.Dense(self.input_dim * 2)
[docs]
def __call__(self, x: jnp.ndarray, cond: jnp.ndarray) -> jnp.ndarray:
"""Applies FiLM modulation.
Parameters
----------
x
Input features of shape (batch, input_dim).
cond
Conditioning features of shape (batch, cond_dim).
Returns
-------
Modulated features of shape (batch, input_dim).
"""
gamma_beta = self.film_generator(cond) # shape: (batch, input_dim * 2)
gamma, beta = jnp.split(gamma_beta, 2, axis=-1) # each shape: (batch, input_dim)
return self.act_fn(gamma * x + beta)
[docs]
class ResNetBlock(nn.Module):
"""Residual conditioning block.
Applies a residual MLP transformation to the input, conditioned on external features.
Parameters
----------
input_dim
Dimensionality of the input features.
projection_dims
Dimensionality of the projection layers.
hidden_dims
Hidden layer sizes for the residual block.
act_fn
Activation function to apply in the MLP block.
dropout_rate
Dropout rate applied after each hidden layer.
"""
input_dim: int
hidden_dims: Sequence[int] = (256, 256)
projection_dims: Sequence[int] = (256, 256)
act_fn: Callable[[jnp.ndarray], jnp.ndarray] = nn.silu
dropout_rate: float = 0.0
[docs]
def setup(self):
self.mlp_block_1 = MLPBlock(dims=self.hidden_dims, act_fn=self.act_fn, dropout_rate=self.dropout_rate)
self.mlp_block_2 = MLPBlock(
dims=list(self.hidden_dims) + [self.input_dim], act_fn=self.act_fn, dropout_rate=self.dropout_rate
)
self.cond_proj = MLPBlock(dims=self.projection_dims, act_fn=self.act_fn, dropout_rate=self.dropout_rate)
[docs]
def __call__(self, x: jnp.ndarray, cond: jnp.ndarray, *, training: bool = True) -> jnp.ndarray:
"""Forward pass of the residual layer.
Parameters
----------
x
Input features of shape (batch, input_dim).
cond
Conditioning features of shape (batch, cond_dim).
training
Whether the model is in training mode.
Returns
-------
Output features of shape (batch, input_dim).
"""
h = self.mlp_block_1(x)
h = h + self.cond_proj(cond)
h = self.mlp_block_2(h)
return h + x
[docs]
class SelfAttention(BaseModule):
"""Self-attention layer
Self-attention layer that can optionally be followed by a FC layer with residual connection,
making it a transformer block.
Parameters
----------
num_heads
Number of heads.
qkv_dim
Dimensionality of the query, key, and value.
dropout_rate
Dropout rate.
transformer_block
Whether to make it a transformer block (adds FC layer with residual connection).
layer_norm
Whether to use layer normalization
"""
num_heads: int = 8
qkv_dim: int = 64
ff_dim: int | None = None
dropout_rate: float = 0.0
transformer_block: bool = False
layer_norm: bool = False
act_fn: Callable[[jnp.ndarray], jnp.ndarray] = nn.silu
[docs]
@nn.compact
def __call__(
self,
x: jnp.ndarray,
mask: jnp.ndarray | None = None,
training: bool = True,
):
"""
Forward pass.
Parameters
----------
x
Input tensor of shape ``(batch_size, set_size, input_dim)`` or
``(batch_size, input_dim)``.
mask
Mask tensor of shape ``(batch_size, 1 | num_heads, set_size, set_size)``.
training
Whether the model is in training mode.
Returns
-------
Output tensor of shape ``(batch_size, set_size, input_dim)``.
"""
squeeze = x.ndim == 2
x = jnp.expand_dims(x, 1) if squeeze else x
# self-attention
z = nn.MultiHeadDotProductAttention(
num_heads=self.num_heads,
qkv_features=self.qkv_dim,
dropout_rate=self.dropout_rate,
)(x, mask=mask, deterministic=not training)
if self.transformer_block:
# query residual connection
z = nn.Dropout(self.dropout_rate)(z, deterministic=not training)
z = z + x
if self.layer_norm:
z = nn.LayerNorm()(z)
# position-wise feed-forward: expand to ff_dim (default 4x input_dim, the canonical
# transformer expansion — NOT qkv_dim), then project back to input_dim.
d_ff = self.ff_dim if self.ff_dim is not None else 4 * z.shape[-1]
z_ = self.act_fn(nn.Dense(d_ff)(z))
z_ = nn.Dense(z.shape[-1])(z_)
z_ = nn.Dropout(self.dropout_rate)(z_, deterministic=not training)
z = z + z_
if self.layer_norm:
z = nn.LayerNorm()(z)
return z.squeeze(1) if squeeze else z
[docs]
@register_layer("self_attention")
class SelfAttentionBlock(BaseModule):
"""
Several self-attention (+ optional FC layer) layers stacked together.
Parameters
----------
num_heads
Number of heads for each layer.
qkv_dim
Dimensionality of the query, key, and value for each layer.
dropout_rate
Dropout rate.
transformer_block
Whether to make layers transformer blocks (adds FC layer with residual connection).
layer_norm
Whether to use layer normalization.
Returns
-------
Output tensor of shape (batch_size, set_size, input_dim).
"""
#: This block consumes an attention mask (see :meth:`__call__`), so
#: :func:`_apply_modules` forwards ``attention_mask`` to it.
takes_attention_mask: ClassVar[bool] = True
num_heads: Sequence[int] | int
qkv_dim: Sequence[int] | int
dropout_rate: float = 0.0
transformer_block: bool = False
layer_norm: bool = False
ff_dim: int | None = None
act_fn: Callable[[jnp.ndarray], jnp.ndarray] = nn.silu
def __post_init__(self) -> None:
"""Initialize the module."""
super().__post_init__()
if not isinstance(self.num_heads, Sequence):
self.num_heads = [self.num_heads]
if not isinstance(self.qkv_dim, Sequence):
self.qkv_dim = [self.qkv_dim]
if len(self.num_heads) != len(self.qkv_dim):
raise ValueError("The number of specified layers should be the same for num_heads and qkv_dims.")
[docs]
@nn.compact
def __call__(
self,
x: jnp.ndarray,
mask: jnp.ndarray | None = None,
training: bool = True,
) -> jnp.ndarray:
"""
Forward pass.
Parameters
----------
x : jnp.ndarray
Input tensor of shape (batch_size, set_size, input_dim).
mask : Optional[jnp.ndarray]
Mask tensor of shape (batch_size, 1 | num_heads, set_size, set_size).
training : bool
Whether the model is in training mode.
Returns
-------
Output tensor of shape (batch_size, set_size, input_dim).
"""
z = x
for num_heads, qkv_dim in zip(self.num_heads, self.qkv_dim, strict=False): # type: ignore[arg-type]
z = SelfAttention(
num_heads=num_heads,
qkv_dim=qkv_dim,
ff_dim=self.ff_dim,
dropout_rate=self.dropout_rate,
transformer_block=self.transformer_block,
layer_norm=self.layer_norm,
act_fn=self.act_fn,
)(z, mask, training)
return z
[docs]
class SeedAttentionPooling(BaseModule):
"""
Pooling by multi-head attention with a trainable seed.
Parameters
----------
num_heads
Number of heads.
v_dim
Dimensionality of the value.
seed_dim
Dimensionality of the seed.
dropout_rate
Dropout rate.
transformer_block
Whether to make it a transformer block (adds FC layer with residual connection).
layer_norm
Whether to use layer normalization.
act_fn
Activation function.
References
----------
:cite:`vaswani:17`
"""
num_heads: int = 8
v_dim: int = 64
seed_dim: int = 64
dropout_rate: float = 0.0
transformer_block: bool = False
layer_norm: bool = False
act_fn: Callable[[jnp.ndarray], jnp.ndarray] = nn.silu
[docs]
@nn.compact
def __call__(
self,
x: jnp.ndarray,
mask: jnp.ndarray | None = None,
training: bool = True,
):
"""
Apply the pooling by multi-head attention.
Parameters
----------
x
Input tensor of shape ``(batch_size, set_size, input_dim)``.
mask
Mask tensor of shape ``(batch_size, 1, set_size, set_size)``.
training
Whether the model is in training mode.
Returns
-------
Output tensor of shape ``(batch_size, input_dim)``.
"""
# trainable seed
S = self.param("S", initializers.xavier_uniform(), (1, 1, self.seed_dim))
S = jnp.tile(S, (x.shape[0], 1, 1))
# multi-head attention
Q = nn.Dense(self.v_dim)(S)
K, V = nn.Dense(self.v_dim)(x), nn.Dense(self.v_dim)(x)
Q_ = jnp.concatenate(jnp.split(Q, self.num_heads, axis=2), axis=0)
K_ = jnp.concatenate(jnp.split(K, self.num_heads, axis=2), axis=0)
V_ = jnp.concatenate(jnp.split(V, self.num_heads, axis=2), axis=0)
A = jnp.matmul(Q_, K_.transpose(0, 2, 1)) / jnp.sqrt(self.v_dim)
if mask is not None:
# mask from (batch_, 1 | num_heads, set_, set_) to (batch_ * num_heads, 1, set_)
mask = jnp.repeat(mask[:, 0, [0], :], self.num_heads, axis=0)
A = jnp.where(mask, A, -1e9)
A = nn.softmax(A)
A = jnp.matmul(A, V_)
if self.transformer_block:
# query residual connection
O = jnp.concatenate(jnp.split(Q_ + A, self.num_heads, axis=0), axis=2)
O = nn.Dropout(rate=self.dropout_rate)(O, deterministic=not training)
if self.layer_norm:
O = nn.LayerNorm()(O)
# FC layer with residual connection
O_ = self.act_fn(nn.Dense(self.v_dim)(O))
O_ = nn.Dropout(rate=self.dropout_rate)(O_, deterministic=not training)
O = O + O_
if self.layer_norm:
O = nn.LayerNorm()(O)
else:
O = jnp.concatenate(jnp.split(A, self.num_heads, axis=0), axis=2)
return O.squeeze(1)
[docs]
class TokenAttentionPooling(BaseModule):
"""
Multi-head attention which aggregates sets by learning a token.
A learnable ``[CLS]`` token is prepended to the set and ``num_layers`` self-attention layers are
applied over ``{CLS, *tokens}`` (BERT/ViT-style); the CLS token's final state is returned as the
aggregate. ``num_layers=1`` with ``transformer_block=False`` reproduces the original single-attention
pooling. Set ``num_layers>=2`` (+ ``transformer_block=True``) to make the condition encoder a
transformer where every token attends to every other before the CLS read-out.
Parameters
----------
num_heads
Number of attention heads (per layer).
qkv_dim
Dimensionality of the query, key, and value (per layer).
dropout_rate
Dropout rate.
num_layers
Number of stacked self-attention layers over ``{CLS, *tokens}``.
transformer_block
Whether each layer is a full transformer block (adds a feed-forward sublayer with residual).
layer_norm
Whether to use layer normalization inside each transformer block.
act_fn
Activation function.
"""
num_heads: int = 8
qkv_dim: int = 64
dropout_rate: float = 0.0
num_layers: int = 1
transformer_block: bool = False
layer_norm: bool = False
ff_dim: int | None = None
act_fn: Callable[[jnp.ndarray], jnp.ndarray] = nn.silu
[docs]
@nn.compact
def __call__(
self,
x: jnp.ndarray,
mask: jnp.ndarray | None = None,
training: bool = True,
) -> jnp.ndarray:
"""Forward pass.
Parameters
----------
x
Input tensor of shape (batch_size, set_size, input_dim).
mask
Mask tensor of shape (batch_size, 1 | num_heads, set_size, set_size).
training
Whether the model is in training mode.
Returns
-------
Output tensor of shape ``(batch_size, input_dim)``.
"""
# add token
token_shape = (len(x), 1)
class_token = nn.Embed(num_embeddings=1, features=x.shape[-1])(jnp.int32(jnp.zeros(token_shape)))
z = jnp.concatenate((class_token, x), axis=-2)
token_mask = jnp.ones((x.shape[0], 1, x.shape[1] + 1, x.shape[1] + 1))
token_mask = token_mask.at[:, :, 1:, 1:].set(mask)
cls_token_to_data = mask[0, 0, :, :].sum(axis=0) > 0
token_mask = token_mask.at[:, :, 0, 1:].set(cls_token_to_data)
token_mask = token_mask.at[:, :, 1:, 0].set(cls_token_to_data)
# transformer over {CLS, *tokens}: num_layers self-attention layers, then read out the CLS.
if self.num_layers == 1 and not self.transformer_block:
# backward-compatible single-attention path: identical module structure (hence Flax param
# tree) to the original TokenAttentionPooling, so historical (pre-refactor) checkpoints
# restore unchanged. num_layers>=2 or transformer_block=True use the SelfAttentionBlock.
z = nn.MultiHeadDotProductAttention(
num_heads=self.num_heads,
qkv_features=self.qkv_dim,
dropout_rate=self.dropout_rate,
)(z, mask=token_mask, deterministic=not training)
else:
z = SelfAttentionBlock(
num_heads=[self.num_heads] * self.num_layers,
qkv_dim=[self.qkv_dim] * self.num_layers,
ff_dim=self.ff_dim,
dropout_rate=self.dropout_rate,
transformer_block=self.transformer_block,
layer_norm=self.layer_norm,
act_fn=self.act_fn,
)(z, token_mask, training)
# only continue with the CLS token (position 0)
return z[:, 0, :]
def _get_layers(
layers: Layers_t,
output_dim: int | None = None,
dropout_rate: float | None = None,
) -> list[nn.Module]:
"""Get modules from layer parameters."""
modules = []
if isinstance(layers, Sequence):
for layer in layers:
layer = dict(layer)
layer_type = layer.pop("layer_type", "mlp")
try:
layer_cls = LAYER_REGISTRY[layer_type]
except KeyError:
raise ValueError(
f"Unknown layer type: {layer_type!r}. Registered layer types: {sorted(LAYER_REGISTRY)}."
) from None
modules.append(layer_cls(**layer))
if output_dim is not None:
modules.append(nn.Dense(output_dim))
if dropout_rate is not None:
modules.append(nn.Dropout(dropout_rate))
return modules
def _apply_modules(
modules: list[nn.Module],
conditions: jax.Array,
attention_mask: jnp.ndarray | None,
training: bool,
) -> jnp.ndarray:
"""Apply modules to conditions."""
for module in modules:
if getattr(module, "takes_attention_mask", False):
conditions = module(conditions, attention_mask, training)
elif isinstance(module, nn.Dense):
conditions = module(conditions)
elif isinstance(module, nn.Dropout):
conditions = module(conditions, deterministic=not training)
else:
conditions = module(conditions, training)
return conditions