[content] utils/nn.py (python)
from functools import partial

from flax import linen as nn
import jax
import jax.numpy as jnp
import einops


class PositionalEncoding(nn.Module):
    """https://uvadlc-notebooks.readthedocs.io/en/latest/tutorial_notebooks/JAX/tutorial6/Transformers_and_MHAttention.html"""

    d_model: int  # Hidden dimensionality of the input.
    max_len: int = 5000  # Maximum length of a sequence to expect.

    def setup(self):
        # Create matrix of [SeqLen, HiddenDim] representing the positional encoding for max_len inputs
        self.pe = jnp.zeros((self.max_len, self.d_model))
        position = jnp.arange(0, self.max_len, dtype=jnp.float32)[:, None]
        div_term = jnp.exp(
            jnp.arange(0, self.d_model, 2) * (-math.log(10000.0) / self.d_model)
        )
        self.pe = self.pe.at[:, 0::2].set(jnp.sin(position * div_term))
        self.pe = self.pe.at[:, 1::2].set(jnp.cos(position * div_term))

    def __call__(self, x):
        x = x + self.pe[: x.shape[2]]
        return x

# class STBlock2(nn.Module):
    # dim: int
    # num_heads: int
    # dropout: float
    # param_dtype: jnp.dtype
    # dtype: jnp.dtype

    # @nn.remat
    # @nn.compact
    # def __call__(self, x: jax.Array) -> jax.Array:
        # # --- Spatial attention ---
        # z = PositionalEncoding(self.dim)(x)
        # z = nn.LayerNorm(
            # param_dtype=self.param_dtype,
            # dtype=self.dtype,
        # )(z)
        # causal_mask = jnp.tri(z.shape[-2])
        # z = nn.MultiHeadAttention(
            # num_heads=self.num_heads,
            # qkv_features=self.dim,
            # dropout_rate=self.dropout,
            # param_dtype=self.param_dtype,
            # dtype=self.dtype,
        # )(z, mask=causal_mask)
        # x = x + z

        # # --- Temporal attention ---
        # x = x.swapaxes(1, 2)
        # z = PositionalEncoding(self.dim)(x)
        # z = nn.LayerNorm(
            # param_dtype=self.param_dtype,
            # dtype=self.dtype,
        # )(z)
        # causal_mask = jnp.tri(z.shape[-2])
        # z = nn.MultiHeadAttention(
            # num_heads=self.num_heads,
            # qkv_features=self.dim,
            # dropout_rate=self.dropout,
            # param_dtype=self.param_dtype,
            # dtype=self.dtype,
        # )(z, mask=causal_mask)
        # x = x + z
        # x = x.swapaxes(1, 2)

        # # --- Feedforward ---
        # z = nn.LayerNorm(
            # param_dtype=self.param_dtype,
            # dtype=self.dtype,
        # )(x)
        # # FIXME (f.srambical): Here, the attention hidden dimension is the same as the FFN's. Usually, FFN hidden dimension is 4x model_dim
        # z = nn.Dense(
            # self.dim,
            # param_dtype=self.param_dtype,
            # dtype=self.dtype,
        # )(z)
        # z = nn.gelu(z)
        # x = x + z

        # return x

# class CausalTransformer(nn.Module):
    # model_dim: int
    # out_dim: int
    # num_blocks: int
    # num_heads: int
    # dropout: float
    # param_dtype: jnp.dtype
    # dtype: jnp.dtype

    # @nn.compact
    # def __call__(self, x: jax.Array) -> jax.Array:
        # # Input projection and normalization
        # x = nn.Sequential(
            # [
                # nn.LayerNorm(
                    # param_dtype=self.param_dtype,
                    # dtype=self.dtype,
                # ),
                # nn.Dense(self.model_dim,
                # param_dtype=self.param_dtype,
                # dtype=self.dtype,
                # ),
                # nn.LayerNorm(
                    # param_dtype=self.param_dtype,
                    # dtype=self.dtype,
                # ),
            # ]
        # )(x)
        # # Causal transformer blocks
        # for _ in range(self.num_blocks):
            # x = STBlock2(
                # dim=self.model_dim,
                # num_heads=self.num_heads,
                # dropout=self.dropout,
                # param_dtype=self.param_dtype,
                # dtype=self.dtype,
            # )(x)

        # # Output projection
        # x = nn.Dense(
            # self.out_dim,
            # param_dtype=self.param_dtype,
            # dtype=self.dtype,
        # )(x)
        # return x  # (B, T, E)


class STBlock(nn.Module):
    dim: int
    num_heads: int
    dropout: float
    param_dtype: jnp.dtype
    dtype: jnp.dtype
    use_flash_attention: bool
    spatial_bert: bool = True

    @nn.remat
    @nn.compact
    def __call__(self, x: jax.Array) -> jax.Array:
        # --- Spatial attention ---
        z = PositionalEncoding(self.dim)(x)
        z = nn.LayerNorm(
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(z)
        z = nn.MultiHeadAttention(
            num_heads=self.num_heads,
            qkv_features=self.dim,
            dropout_rate=self.dropout,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
            attention_fn=_create_flash_attention_fn(self.use_flash_attention, is_causal=not self.spatial_bert),
        )(z)
        x = x + z

        # --- Temporal attention ---
        x = x.swapaxes(1, 2)
        z = PositionalEncoding(self.dim)(x)
        z = nn.LayerNorm(
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(z)
        causal_mask = jnp.tri(z.shape[-2])
        z = nn.MultiHeadAttention(
            num_heads=self.num_heads,
            qkv_features=self.dim,
            dropout_rate=self.dropout,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
            attention_fn=_create_flash_attention_fn(self.use_flash_attention, is_causal=True),
        # FIXME (f.srambical): check whether we should still pass the mask if we set is_causal=True
        )(z)
        x = x + z
        x = x.swapaxes(1, 2)

        # --- Feedforward ---
        z = nn.LayerNorm(
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(x)
        # FIXME (f.srambical): Here, the attention hidden dimension is the same as the FFN's. Usually, FFN hidden dimension is 4x model_dim
        z = nn.Dense(
            self.dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(z)
        z = nn.gelu(z)
        x = x + z

        return x


class STTransformer(nn.Module):
    model_dim: int
    out_dim: int
    num_blocks: int
    num_heads: int
    dropout: float
    param_dtype: jnp.dtype
    dtype: jnp.dtype
    use_flash_attention: bool
    spatial_bert: bool = True

    @nn.compact
    def __call__(self, x: jax.Array) -> jax.Array:
        x = nn.Sequential(
            [
                nn.LayerNorm(
                    param_dtype=self.param_dtype,
                    dtype=self.dtype,
                ),
                nn.Dense(self.model_dim,
                param_dtype=self.param_dtype,
                dtype=self.dtype,
                ),
                nn.LayerNorm(
                    param_dtype=self.param_dtype,
                    dtype=self.dtype,
                ),
            ]
        )(x)
        for _ in range(self.num_blocks):
            x = STBlock(
                dim=self.model_dim,
                num_heads=self.num_heads,
                dropout=self.dropout,
                param_dtype=self.param_dtype,
                dtype=self.dtype,
                spatial_bert=self.spatial_bert,