[tab] sampling.py (python)
from abc import ABC, abstractmethod
from typing import Dict, Protocol
import jax
import jax.numpy as jnp
import flax.nnx as nnx


class DynamicsModel(Protocol):
    """Protocol defining the interface that Dynamics models must implement for sampling."""
    
    def tokenizer(self) -> nnx.Module:
        """Return the tokenizer component."""
        ...
    
    def lam(self) -> nnx.Module:
        """Return the LAM component."""
        ...
    
    def dynamics(self) -> nnx.Module:
        """Return the dynamics component."""
        ...


class BaseSampler(ABC):
    """Abstract base class for sampling strategies."""
    
    def __init__(self, model: DynamicsModel):
        self.model = model
    
    @abstractmethod
    def sample(
        self,
        batch: Dict[str, jax.Array],
        seq_len: int,
        temperature: float = 1.0,
        sample_argmax: bool = False,
        **kwargs
    ) -> jax.Array:
        """Sample future frames given input context."""
        pass


class MaskGITSampler(BaseSampler):
    """MaskGIT sampling strategy."""
    
    def sample(
        self,
        batch: Dict[str, jax.Array],
        seq_len: int,
        steps: int = 25,
        temperature: float = 1.0,
        sample_argmax: bool = False,
    ) -> jax.Array:
        """
        Autoregressively samples up to `seq_len` future frames using MaskGIT strategy.
        
        - Input frames are tokenized once.
        - Future frames are generated autoregressively in token space.
        - All frames are detokenized in a single pass.
        """
        # --- Encode videos and actions ---
        tokenizer_out = self.model.tokenizer.vq_encode(batch["videos"], training=False)
        token_idxs = tokenizer_out["indices"]  # (B, T, N)
        B, T, N = token_idxs.shape
        pad_shape = (B, seq_len - T, N)
        pad = jnp.zeros(pad_shape, dtype=token_idxs.dtype)
        token_idxs = jnp.concatenate([token_idxs, pad], axis=1)  # (B, S, N)
        action_tokens = self.model.lam.vq.get_codes(batch["latent_actions"])

        def maskgit_step_fn(
            carry: tuple[jax.Array, jax.Array, jax.Array, jax.Array], step: jax.Array
        ) -> tuple[tuple[jax.Array, jax.Array, jax.Array, jax.Array], None]:
            rng, token_idxs, mask, action_tokens = carry
            N = token_idxs.shape[2]

            # --- Construct + encode video ---
            vid_embed = self.model.dynamics.patch_embed(token_idxs)  # (B, S, N, D)
            mask_token = self.model.dynamics.mask_token.value  # (1, 1, 1, D,)
            mask_expanded = mask[..., None]  # (B, S, N, 1)
            vid_embed = jnp.where(mask_expanded, mask_token, vid_embed)

            # --- Predict transition ---
            act_embed = self.model.dynamics.action_up(action_tokens)
            vid_embed += jnp.pad(act_embed, ((0, 0), (1, 0), (0, 0), (0, 0)))
            unmasked_ratio = jnp.cos(jnp.pi * (step + 1) / (steps * 2))
            step_temp = temperature * (1.0 - unmasked_ratio)
            final_logits = self.model.dynamics.dynamics(vid_embed) / step_temp

            # --- Sample new tokens for final frame ---
            if sample_argmax:
                sampled_token_idxs = jnp.argmax(final_logits, axis=-1)
            else:
                rng, _rng = jax.random.split(rng)
                sampled_token_idxs = jax.random.categorical(_rng, final_logits)
            gather_fn = jax.vmap(jax.vmap(jax.vmap(lambda x, y: x[y])))
            final_token_probs = gather_fn(
                jax.nn.softmax(final_logits), sampled_token_idxs
            )
            final_token_probs += ~mask
            # Update masked tokens only
            token_idxs = jnp.where(mask, sampled_token_idxs, token_idxs)

            # --- Update mask ---
            num_unmasked_tokens = jnp.round(N * (1.0 - unmasked_ratio)).astype(int)
            idx_mask = jnp.arange(final_token_probs.shape[-1]) > num_unmasked_tokens
            sorted_idxs = jnp.argsort(final_token_probs, axis=-1, descending=True)
            mask_update_fn = jax.vmap(lambda msk, ids: msk.at[ids].set(idx_mask))
            new_mask = mask_update_fn(mask, sorted_idxs)

            new_carry = (rng, token_idxs, new_mask, action_tokens)
            return new_carry, None

        def generation_step_fn(
            carry: tuple[jax.Array, jax.Array], step_t: jax.Array
        ) -> tuple[tuple[jax.Array, jax.Array], None]:
            rng, current_token_idxs = carry
            rng, step_rng = jax.random.split(rng)

            # Mask current and future frames (i.e., t >= step_t)
            mask = jnp.arange(seq_len) >= step_t  # (S,)
            mask = jnp.broadcast_to(mask[None, :, None], (B, seq_len, N)).astype(
                bool
            )  # (B, S, N)
            masked_token_idxs = current_token_idxs * ~mask

            # --- Initialize and run MaskGIT loop ---
            init_carry_maskgit = (
                step_rng,
                masked_token_idxs,
                mask,
                action_tokens,
            )
            final_carry_maskgit, _ = jax.lax.scan(
                maskgit_step_fn, init_carry_maskgit, jnp.arange(steps)
            )
            updated_token_idxs = final_carry_maskgit[1]
            new_carry = (rng, updated_token_idxs)
            return new_carry, None

        # --- Run the autoregressive generation using jax.lax.scan ---
        initial_carry = (batch["rng"], token_idxs)
        timesteps_to_scan = jnp.arange(T, seq_len)
        final_carry, _ = jax.lax.scan(
            generation_step_fn, initial_carry, timesteps_to_scan
        )
        final_token_idxs = final_carry[1]

        # --- Decode all tokens at once at the end ---
        H, W = batch["videos"].shape[2:4]
        final_frames = self.model.tokenizer.decode(
            final_token_idxs,
            video_hw=(H, W),
        )
        return final_frames


class CausalSampler(BaseSampler):
    """Causal transformer sampling strategy."""
    
    def sample(
        self,
        batch: Dict[str, jax.Array],
        seq_len: int,
        temperature: float = 1.0,
        sample_argmax: bool = False,
    ) -> jax.Array:
        """
        Autoregressively samples up to `seq_len` future frames using causal transformer.
        
        - Input frames are tokenized once.
        - Future frames are generated one at a time, each conditioned on all previous frames.
        - All frames are detokenized in a single pass at the end.
        """
        # --- Encode context frames ---
        tokenizer_out = self.model.tokenizer.vq_encode(batch["videos"], training=False)
        token_idxs = tokenizer_out["indices"]  # (B, T, N)
        B, T, N = token_idxs.shape

        # --- Prepare initial token sequence ---
        # Pad with zeros for future frames
        pad_shape = (B, seq_len - T, N)
        token_idxs_full = jnp.concatenate(
            [token_idxs, jnp.zeros(pad_shape, dtype=token_idxs.dtype)], axis=1
        )  # (B, seq_len, N)

        # --- Prepare latent actions ---
        action_tokens = self.model.lam.vq.get_codes(batch["latent_actions"])

        def token_step_fn(
            carry: tuple[jax.Array, jax.Array, jax.Array], token_idx: jax.Array
        ) -> tuple[tuple[jax.Array, jax.Array, jax.Array], None]:
            rng, token_idxs_full, action_tokens = carry
            t = token_idx // N
            n = token_idx % N
            
            dyna_inputs = {
                "video_tokens": token_idxs_full,
                "latent_actions": action_tokens,
            }
            next_token_logits, _ = self.model.dynamics(dyna_inputs, training=False)
            next_token_logits = next_token_logits[:, t, n, :].astype(
                jnp.float32
            )  # (B, vocab_size)

            if sample_argmax:
                next_token = jnp.argmax(next_token_logits, axis=-1)  # (B,)
            else:
                rng, step_rng = jax.random.split(rng)
                next_token = jax.random.categorical(
                    step_rng, next_token_logits / temperature, axis=-1
                )  # (B,)

            # Insert the generated tokens into the sequence
            token_idxs_full = token_idxs_full.at[:, t, n].set(next_token)
            
            new_carry = (rng, token_idxs_full, action_tokens)
            return new_carry, None

        # --- Autoregressive generation ---
        future_frames = seq_len - T
        total_future_tokens = future_frames * N
        start_token_idx = T * N
        step_indices = jnp.arange(start_token_idx, start_token_idx + total_future_tokens)
        
        initial_carry = (batch["rng"], token_idxs_full, action_tokens)
        final_carry, _ = jax.lax.scan(
            token_step_fn, initial_carry, step_indices
        )
        final_token_idxs = final_carry[1]

        # --- Decode all tokens at once at the end ---
        H, W = batch["videos"].shape[2:4]
        final_frames = self.model.tokenizer.decode(final_token_idxs, video_hw=(H, W))
        return final_frames


def create_sampler(dynamics_type: str, model: DynamicsModel) -> BaseSampler:
    """Factory function to create the appropriate sampler based on dynamics type."""
    if dynamics_type == "maskgit":
        return MaskGITSampler(model)
    elif dynamics_type == "causal":
        return CausalSampler(model)
    else:
        raise ValueError(f"Unknown dynamics type: {dynamics_type}")