[content] jasmine/utils/nn.py (python)
class AxialBlock(nnx.Module):
    """Axial transformer block"""

    def __init__(
        self,
        dim: int,
        ffn_dim: int,
        num_heads: int,
        dropout: float,
        param_dtype: jnp.dtype,
        dtype: jnp.dtype,
        use_flash_attention: bool,
        rngs: nnx.Rngs,
        spatial_causal: bool,
        temporal_causal: bool,
        sow_weights: bool,
        sow_activations: bool,
        decode: bool,
    ):
        self.dim = dim
        self.ffn_dim = ffn_dim
        self.num_heads = num_heads
        self.dropout = dropout
        self.param_dtype = param_dtype
        self.dtype = dtype
        self.use_flash_attention = use_flash_attention
        self.spatial_causal = spatial_causal
        self.temporal_causal = temporal_causal
        self.sow_weights = sow_weights
        self.sow_activations = sow_activations
        self.decode = decode
        self.spatial_norm = nnx.LayerNorm(
            num_features=self.dim,
            param_dtype=self.param_dtype,
            dtype=self.param_dtype,  # layer norm in full precision
            rngs=rngs,
        )
        self.spatial_attention = nnx.MultiHeadAttention(
            num_heads=self.num_heads,
            in_features=self.dim,
            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=self.spatial_causal
            ),
            rngs=rngs,
            decode=self.decode,
        )

        self.temporal_norm = nnx.LayerNorm(
            num_features=self.dim,
            param_dtype=self.param_dtype,
            dtype=self.param_dtype,  # layer norm in full precision
            rngs=rngs,
        )
        self.temporal_attention = nnx.MultiHeadAttention(
            num_heads=self.num_heads,
            in_features=self.dim,
            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=self.temporal_causal
            ),
            rngs=rngs,
            decode=self.decode,
        )

        self.ffn_norm = nnx.LayerNorm(
            num_features=self.dim,
            param_dtype=self.param_dtype,
            dtype=self.param_dtype,  # layer norm in full precision
            rngs=rngs,
        )
        self.ffn_dense1 = nnx.Linear(
            in_features=self.dim,
            out_features=self.ffn_dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
            rngs=rngs,
        )
        self.ffn_dense2 = nnx.Linear(
            in_features=self.ffn_dim,
            out_features=self.dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
            rngs=rngs,
        )

    @nnx.remat
    def __call__(self, x_BTNM: jax.Array) -> jax.Array:
        # --- Spatial attention ---
        z_BTNM = self.spatial_norm(x_BTNM)
        z_BTNM = self.spatial_attention(z_BTNM, sow_weights=self.sow_weights)
        x_BTNM = x_BTNM + z_BTNM

        # --- Temporal attention ---
        x_BNTM = x_BTNM.swapaxes(1, 2)
        z_BNTM = self.temporal_norm(x_BNTM)
        z_BNTM = self.temporal_attention(z_BNTM, sow_weights=self.sow_weights)
        x_BNTM = x_BNTM + z_BNTM
        x_BTNM = x_BNTM.swapaxes(1, 2)

        # --- Feedforward ---
        z_BTNM = self.ffn_norm(x_BTNM)
        z_BTND = self.ffn_dense1(z_BTNM)
        z_BTND = jax.nn.gelu(z_BTND)
        z_BTNM = self.ffn_dense2(z_BTND)
        x_BTNM = x_BTNM + z_BTNM
        if self.sow_activations:
            self.sow(nnx.Intermediate, "activations", x_BTNM)
        return x_BTNM


class AxialTransformer(nnx.Module):
    """
    Axial transformer

    Dimension keys:
        B: batch size
        T: number of frames
        N: number of patches per frame
        I: number of input features
        M: model dimension
        D: FFN dimension
        O: output dimension
    """

    def __init__(
        self,
        input_dim: int,
        model_dim: int,
        ffn_dim: int,
        out_dim: int,
        num_blocks: int,
        num_heads: int,
        dropout: float,
        param_dtype: jnp.dtype,
        dtype: jnp.dtype,
        use_flash_attention: bool,
        rngs: nnx.Rngs,
        spatial_causal: bool,
        temporal_causal: bool,
        decode: bool = False,
        sow_weights: bool = False,
        sow_activations: bool = False,
        sow_logits: bool = False,
        max_len: int = 5000,
    ):
        self.input_dim = input_dim
        self.model_dim = model_dim
        self.ffn_dim = ffn_dim
        self.out_dim = out_dim
        self.num_blocks = num_blocks
        self.num_heads = num_heads
        self.dropout = dropout
        self.param_dtype = param_dtype
        self.dtype = dtype
        self.use_flash_attention = use_flash_attention
        self.spatial_causal = spatial_causal
        self.temporal_causal = temporal_causal
        self.sow_logits = sow_logits
        self.sow_weights = sow_weights
        self.sow_activations = sow_activations
        self.decode = decode
        self.input_norm1 = nnx.LayerNorm(
            num_features=self.input_dim,
            param_dtype=self.param_dtype,
            dtype=self.param_dtype,  # layer norm in full precision
            rngs=rngs,
        )
        self.input_dense = nnx.Linear(
            in_features=self.input_dim,
            out_features=self.model_dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
            rngs=rngs,
        )
        self.input_norm2 = nnx.LayerNorm(
            num_features=self.model_dim,
            param_dtype=self.param_dtype,
            dtype=self.param_dtype,  # layer norm in full precision
            rngs=rngs,
        )

        self.pos_enc = _get_spatiotemporal_positional_encoding(
            self.model_dim, max_len=max_len
        )

        self.blocks = []
        for _ in range(self.num_blocks):
            self.blocks.append(
                AxialBlock(
                    dim=self.model_dim,
                    ffn_dim=self.ffn_dim,
                    num_heads=self.num_heads,
                    dropout=self.dropout,
                    param_dtype=self.param_dtype,
                    dtype=self.dtype,
                    use_flash_attention=self.use_flash_attention,
                    rngs=rngs,
                    spatial_causal=self.spatial_causal,
                    temporal_causal=self.temporal_causal,
                    sow_weights=self.sow_weights,
                    sow_activations=self.sow_activations,
                    decode=self.decode,
                )
            )

        self.output_dense = nnx.Linear(
            in_features=self.model_dim,
            out_features=self.out_dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
            rngs=rngs,
        )

    def __call__(self, x_BTNI: jax.Array) -> jax.Array:
        x_BTNI = self.input_norm1(x_BTNI)
        x_BTNM = self.input_dense(x_BTNI)
        x_BTNM = self.input_norm2(x_BTNM)
        x_BTNM = self.pos_enc(x_BTNM)
        for block in self.blocks:
            x_BTNM = block(x_BTNM)

        x_BTNV = self.output_dense(x_BTNM)
        if self.sow_logits:
            self.sow(nnx.Intermediate, "logits", x_BTNV)
        return x_BTNV