[content] models/dynamics.py (python)
from typing import Dict, Any

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

from utils.nn import STTransformer


class DynamicsMaskGIT(nn.Module):
    """MaskGIT dynamics model"""

    model_dim: int
    num_latents: int
    num_blocks: int
    num_heads: int
    dropout: float
    mask_limit: float
    param_dtype: jnp.dtype
    dtype: jnp.dtype
    use_flash_attention: bool

    def setup(self):
        self.dynamics = STTransformer(
            model_dim=self.model_dim,
            out_dim=self.num_latents,
            num_blocks=self.num_blocks,
            num_heads=self.num_heads,
            dropout=self.dropout,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
            use_flash_attention=self.use_flash_attention,
            use_kv_cache=False,
            spatial_bert=True,
        )
        self.patch_embed = nn.Embed(self.num_latents, self.model_dim)
        self.mask_token = self.param(
            "mask_token",
            nn.initializers.lecun_uniform(),
            (1, 1, 1, self.model_dim),
        )
        self.action_up = nn.Dense(
            self.model_dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )

    def init_cache(self, input_shape, dtype):
        """
        Initialize the cache for autoregressive generation.
        
        Args:
            input_shape: Shape of the input (B, seq_len, N, model_dim)
            dtype: Data type for the cache
            
        Returns:
            Cache structure for the STTransformer attention layers
        """
        B, seq_len, N, model_dim = input_shape
        head_dim = model_dim // self.num_heads
        
        # Create cache structure for each STBlock
        cache = {}
        for i in range(self.num_blocks):
            block_cache = {}
            
            # Spatial attention cache (MultiHeadAttention_0)
            # Shape: (B, seq_len, num_heads, head_dim)
            spatial_key_shape = (B, seq_len, self.num_heads, head_dim)
            spatial_value_shape = (B, seq_len, self.num_heads, head_dim)
            
            block_cache['MultiHeadAttention_0'] = {
                'cached_key': jnp.zeros(spatial_key_shape, dtype=dtype),
                'cached_value': jnp.zeros(spatial_value_shape, dtype=dtype),
                'cache_index': jnp.array(0, dtype=jnp.int32)
            }
            
            # Temporal attention cache (MultiHeadAttention_1)
            # For temporal attention, we need to swap axes: (B, N, seq_len, model_dim)
            temporal_key_shape = (B, N, seq_len, self.num_heads, head_dim)
            temporal_value_shape = (B, N, seq_len, self.num_heads, head_dim)
            
            block_cache['MultiHeadAttention_1'] = {
                'cached_key': jnp.zeros(temporal_key_shape, dtype=dtype),
                'cached_value': jnp.zeros(temporal_value_shape, dtype=dtype),
                'cache_index': jnp.array(0, dtype=jnp.int32)
            }
            
            cache[f'STBlock_{i}'] = block_cache
            
        return cache

    def __call__(self, batch: Dict[str, Any], training: bool = True) -> Dict[str, Any]:
        # --- Mask videos ---
        vid_embed = self.patch_embed(batch["video_tokens"])
        if training:
            rng1, rng2 = jax.random.split(batch["mask_rng"])
            mask_prob = jax.random.uniform(rng1, minval=self.mask_limit)
            mask = jax.random.bernoulli(rng2, mask_prob, vid_embed.shape[:-1])
            mask = mask.at[:, 0].set(False)
            vid_embed = jnp.where(jnp.expand_dims(mask, -1), self.mask_token, vid_embed)
        else:
            mask = None

        # --- Predict transition ---
        act_embed = self.action_up(batch["latent_actions"])
        vid_embed += jnp.pad(act_embed, ((0, 0), (1, 0), (0, 0), (0, 0)))
        logits = self.dynamics(vid_embed)
        return dict(token_logits=logits, mask=mask)


class DynamicsAutoregressive(nn.Module):
    """Autoregressive (causal) dynamics model"""

    model_dim: int
    num_latents: int
    num_blocks: int
    num_heads: int
    dropout: float
    param_dtype: jnp.dtype
    dtype: jnp.dtype
    use_flash_attention: bool
    is_training: bool

    def setup(self):
        self.dynamics = STTransformer(
            model_dim=self.model_dim,
            out_dim=self.num_latents,
            num_blocks=self.num_blocks,
            num_heads=self.num_heads,
            dropout=self.dropout,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
            use_flash_attention=self.use_flash_attention,
            use_kv_cache=not self.is_training,
            spatial_bert=False,
        )
        self.patch_embed = nn.Embed(self.num_latents, self.model_dim)
        self.action_up = nn.Dense(
            self.model_dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )

    def init_cache(self, input_shape, dtype):
        """
        Initialize the cache for autoregressive generation.
        
        Args:
            input_shape: Shape of the input (B, seq_len, N, model_dim)
            dtype: Data type for the cache
            
        Returns:
            Cache structure for the STTransformer attention layers
        """
        B, seq_len, N, model_dim = input_shape
        head_dim = model_dim // self.num_heads
        
        # Create cache structure for each STBlock
        cache = {}
        for i in range(self.num_blocks):
            block_cache = {}
            
            # Spatial attention cache (MultiHeadAttention_0)
            # Shape: (B, seq_len, num_heads, head_dim)
            spatial_key_shape = (B, seq_len, self.num_heads, head_dim)
            spatial_value_shape = (B, seq_len, self.num_heads, head_dim)
            
            block_cache['MultiHeadAttention_0'] = {
                'cached_key': jnp.zeros(spatial_key_shape, dtype=dtype),
                'cached_value': jnp.zeros(spatial_value_shape, dtype=dtype),
                'cache_index': jnp.array(0, dtype=jnp.int32)
            }
            
            # Temporal attention cache (MultiHeadAttention_1)
            # For temporal attention, we need to swap axes: (B, N, seq_len, model_dim)
            temporal_key_shape = (B, N, seq_len, self.num_heads, head_dim)
            temporal_value_shape = (B, N, seq_len, self.num_heads, head_dim)
            
            block_cache['MultiHeadAttention_1'] = {
                'cached_key': jnp.zeros(temporal_key_shape, dtype=dtype),
                'cached_value': jnp.zeros(temporal_value_shape, dtype=dtype),
                'cache_index': jnp.array(0, dtype=jnp.int32)
            }
            
            cache[f'STBlock_{i}'] = block_cache
            
        return cache

    def __call__(self, batch: Dict[str, Any], training: bool = True) -> Dict[str, Any]:
        vid_embed = self.patch_embed(batch["video_tokens"])
        act_embed = self.action_up(batch["latent_actions"])
        vid_embed += jnp.pad(act_embed, ((0, 0), (1, 0), (0, 0), (0, 0)))
        # vid_embed_padded = jnp.pad(vid_embed, ((0, 0), (1, 0), (1, 0), (0, 0)))
        # logits = self.dynamics(vid_embed_padded)[:, :-1, :-1]
        vid_embed_padded = jnp.pad(vid_embed, ((0, 0), (0, 0), (1, 0), (0, 0)))
        logits = self.dynamics(vid_embed_padded)[:, :, :-1]
        mask = jnp.ones(vid_embed.shape[:-1])
        # next_tokens = jnp.argmax(logits, axis=-1)
        # print(next_tokens.shape)
        # jax.debug.breakpoint()
        return dict(token_logits=logits, mask=mask)