[tab] genie.py (python)
from typing import Dict, Any

import optax
import jax
import jax.numpy as jnp
import flax.nnx as nnx
from flax.training.train_state import TrainState
import orbax.checkpoint as ocp

from models.dynamics import DynamicsMaskGIT
from models.lam import LatentActionModel
from models.tokenizer import TokenizerVQVAE

import grain


class Genie(nnx.Module):
    """Genie model"""

    def __init__(self, in_dim: int, tokenizer_dim: int, tokenizer_ffn_dim: int, 
                 latent_patch_dim: int, num_patch_latents: int, patch_size: int, 
                 tokenizer_num_blocks: int, tokenizer_num_heads: int, lam_dim: int, 
                 lam_ffn_dim: int, latent_action_dim: int, num_latent_actions: int, 
                 lam_patch_size: int, lam_num_blocks: int, lam_num_heads: int, 
                 lam_co_train: bool, dyna_dim: int, dyna_ffn_dim: int, 
                 dyna_num_blocks: int, dyna_num_heads: int, param_dtype: jnp.dtype, 
                 dtype: jnp.dtype, use_flash_attention: bool, rngs: nnx.Rngs, 
                 dropout: float = 0.0, mask_limit: float = 0.0):
        # --- Tokenizer ---
        self.in_dim = in_dim
        self.tokenizer_dim = tokenizer_dim
        self.tokenizer_ffn_dim = tokenizer_ffn_dim
        self.latent_patch_dim = latent_patch_dim
        self.num_patch_latents = num_patch_latents
        self.patch_size = patch_size
        self.tokenizer_num_blocks = tokenizer_num_blocks
        self.tokenizer_num_heads = tokenizer_num_heads
        # --- LAM ---
        self.lam_dim = lam_dim
        self.lam_ffn_dim = lam_ffn_dim
        self.latent_action_dim = latent_action_dim
        self.num_latent_actions = num_latent_actions
        self.lam_patch_size = lam_patch_size
        self.lam_num_blocks = lam_num_blocks
        self.lam_num_heads = lam_num_heads
        self.lam_co_train = lam_co_train
        # --- Dynamics ---
        self.dyna_dim = dyna_dim
        self.dyna_ffn_dim = dyna_ffn_dim
        self.dyna_num_blocks = dyna_num_blocks
        self.dyna_num_heads = dyna_num_heads
        self.param_dtype = param_dtype
        self.dtype = dtype
        self.use_flash_attention = use_flash_attention
        self.dropout = dropout
        self.mask_limit = mask_limit
        
        self.tokenizer = TokenizerVQVAE(
            in_dim=self.in_dim,
            model_dim=self.tokenizer_dim,
            ffn_dim=self.tokenizer_ffn_dim,
            latent_dim=self.latent_patch_dim,
            num_latents=self.num_patch_latents,
            patch_size=self.patch_size,
            num_blocks=self.tokenizer_num_blocks,
            num_heads=self.tokenizer_num_heads,
            dropout=0.0,
            codebook_dropout=0.0,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
            use_flash_attention=self.use_flash_attention,
            rngs=rngs,
        )
        self.lam = LatentActionModel(
            in_dim=self.in_dim,
            model_dim=self.lam_dim,
            ffn_dim=self.lam_ffn_dim,
            latent_dim=self.latent_patch_dim,
            num_latents=self.num_latent_actions,
            patch_size=self.lam_patch_size,
            num_blocks=self.lam_num_blocks,
            num_heads=self.lam_num_heads,
            dropout=0.0,
            codebook_dropout=0.0,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
            use_flash_attention=self.use_flash_attention,
            rngs=rngs,
        )
        self.dynamics = DynamicsMaskGIT(
            model_dim=self.dyna_dim,
            ffn_dim=self.dyna_ffn_dim,
            num_latents=self.num_patch_latents,
            num_blocks=self.dyna_num_blocks,
            num_heads=self.dyna_num_heads,
            dropout=self.dropout,
            mask_limit=self.mask_limit,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
            use_flash_attention=self.use_flash_attention,
            rngs=rngs,
        )

    def __call__(self, batch: Dict[str, Any], training: bool = True) -> Dict[str, Any]:
        tokenizer_outputs = self.tokenizer.vq_encode(batch["videos"], training=False)
        lam_outputs = self.lam.vq_encode(batch["videos"], training=False)
        latent_actions = jax.lax.cond(
            self.lam_co_train,
            lambda: lam_outputs["z_q"],
            lambda: jax.lax.stop_gradient(lam_outputs["z_q"])
        )
        outputs = dict(
            video_tokens=jax.lax.stop_gradient(tokenizer_outputs["indices"]),
            latent_actions=latent_actions,
        )
        outputs["mask_rng"] = batch["mask_rng"]
        dyna_outputs = self.dynamics(outputs, training)
        outputs.update(dyna_outputs)
        mle_indices = jnp.argmax(outputs["token_logits"], axis=-1)
        outputs["recon"] = self.tokenizer.decode(
            mle_indices, batch["videos"].shape[2:4]
        )
        outputs["lam_indices"] = lam_outputs["indices"]
        return outputs

    def sample(
        self,
        batch: Dict[str, Any],
        seq_len: int,
        steps: int = 25,
        temperature: float = 1,
        sample_argmax: bool = False,
    ) -> Any:
        """
        Autoregressively samples up to `seq_len` future frames, following Figure 8 of the paper.

        - Input frames are tokenized once.
        - Future frames are generated autoregressively in token space.
        - All frames are detokenized in a single pass.

        Note:
        - For interactive or step-wise sampling, detokenization should occur after each action.
        - To maintain consistent tensor shapes across timesteps, all current and future frames are decoded at every step.
        - Temporal causal structure is preserved by 
            a) reapplying the mask before each decoding step.
            b) a temporal causal mask is applied within each ST-transformer block.

        Dimension keys:
            B: batch size  
            T: number of input (conditioning) frames  
            N: patches per frame  
            S: sequence length  
            A: action space  
            D: model latent dimension
        """
        # --- Encode videos and actions ---
        tokenizer_out = self.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.lam.vq.get_codes(batch["latent_actions"])

        maskgit_step = MaskGITStep(
            dynamics=self.dynamics,
            tokenizer=self.tokenizer,
            temperature=temperature,
            sample_argmax=sample_argmax,
            steps=steps,
        )

        def maskgit_loop_fn(carry, module, x):
            return module(carry, x)
        
        scanned_maskgit_loop = nnx.scan(
            maskgit_loop_fn,
            in_axes=(nnx.Carry, None, 0),
            out_axes=(nnx.Carry, 0),
            length=steps
        )

        def generation_step_fn(carry, step_t):
            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)) # (B, S, N)
            mask = mask.astype(bool)
            masked_token_idxs = current_token_idxs * ~mask

            # --- Initialize and run MaskGIT loop ---
            init_carry_maskgit = (
                step_rng,
                masked_token_idxs,
                mask,
                action_tokens,
            )
            # FIXME (f.srambical): test whether sampling works with this
            final_carry_maskgit, _ = scanned_maskgit_loop(
                init_carry_maskgit, maskgit_step, jnp.arange(steps)
            )
            final_carry_maskgit = carry
            updated_token_idxs = final_carry_maskgit[1]
            new_carry = (rng, updated_token_idxs)
            return new_carry, None

        # --- Run the autoregressive generation using 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 ---
        final_frames = self.tokenizer.decode(
            final_token_idxs,
            video_hw=batch["videos"].shape[2:4],
        )
        return final_frames

    def vq_encode(self, batch, training) -> Dict[str, Any]:
        # --- Preprocess videos ---
        lam_output = self.lam.vq_encode(batch["videos"], training=training)
        return lam_output["indices"]


class MaskGITStep(nnx.Module):
    def __init__(self, dynamics: DynamicsMaskGIT, tokenizer: TokenizerVQVAE, 
                 temperature: float, sample_argmax: bool, steps: int):
        self.dynamics = dynamics
        self.tokenizer = tokenizer
        self.temperature = temperature
        self.sample_argmax = sample_argmax
        self.steps = steps

    def __call__(self, carry, x):
        rng, token_idxs, mask, action_tokens = carry
        step = x
        N = token_idxs.shape[2]

        # --- Construct + encode video ---
        vid_embed = self.dynamics.patch_embed(token_idxs) # (B, S, N, D)
        mask_token = self.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.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) / (self.steps * 2))
        step_temp = self.temperature * (1.0 - unmasked_ratio)
        final_logits = self.dynamics.dynamics(vid_embed) / step_temp

        # --- Sample new tokens for final frame ---
        if self.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

# FIXME (f.srambical): adjust restore_genie_components to work with new NNX checkpointing logic
# FIXME (f.srambical): add conversion script for old checkpoints
def restore_genie_components(
    train_state: TrainState,
    sharding: jax.sharding.NamedSharding,
    grain_iterator: grain.DataLoaderIterator,
    inputs: Dict[str, jax.Array],
    rng: jax.Array,
    args,
):
    """Restore pre-trained Genie components"""
    pass

def _create_abstract_sharded_pytree(pytree_template, sharding_spec):
    """Replaces arrays in a pytree with ShapeDtypeStructs having the given sharding."""

    def map_fn(leaf_template):
        if hasattr(leaf_template, "shape") and hasattr(leaf_template, "dtype"):
            return jax.ShapeDtypeStruct(
                leaf_template.shape, leaf_template.dtype, sharding=sharding_spec
            )
        return leaf_template

    return jax.tree_util.tree_map(map_fn, pytree_template)