[content] utils/nn.py (python)

class BidirectionalSTBlock(nn.Module):
    dim: int
    ffn_dim: int
    num_heads: int
    dropout: float
    param_dtype: jnp.dtype
    dtype: jnp.dtype
    use_flash_attention: bool

    @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=False
            ),
        )(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, mask=causal_mask)
        x = x + z
        x = x.swapaxes(1, 2)

        # --- Feedforward ---
        z = nn.LayerNorm(
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(x)
        z = nn.Dense(
            self.ffn_dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(z)
        z = nn.gelu(z)
        z = nn.Dense(
            self.dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(z)
        x = x + z

        return x


class BidirectionalSTTransformer(nn.Module):
    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

    @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 = BidirectionalSTBlock(
                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,
            )(x)
        x = nn.Dense(
            self.out_dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(x)
        return x  # (B, T, E)