p-doom/crowd-code-dataset-0.2 huggingface dataset:
[tab] models/dynamics.py (python)
from typing import Dict, Any

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

from utils.nn import STTransformer


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

    model_dim: int
    ffn_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(
            self.model_dim,
            self.ffn_dim,
            self.num_latents,
            self.num_blocks,
            self.num_heads,
            self.dropout,
            self.param_dtype,
            self.dtype,
            use_flash_attention=self.use_flash_attention,
        )
        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 __call__(self, batch: Dict[str, Any], training: bool = True) -> Dict[str, Any]:
        # --- Mask videos ---
        vid_embed = self.patch_embed(batch["video_tokens"])
        if training:
            batch_size = vid_embed.shape[0]
            _rng_prob, *_rngs_mask = jax.random.split(batch["mask_rng"], batch_size + 1)
            mask_prob = jax.random.uniform(
                _rng_prob, shape=(batch_size,), minval=self.mask_limit
            )
            per_sample_shape = vid_embed.shape[1:-1]
            mask = jax.vmap(
                lambda rng, prob: jax.random.bernoulli(rng, prob, per_sample_shape),
                in_axes=(0, 0),
            )(jnp.asarray(_rngs_mask), mask_prob)
            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)
[content] models/dynamics.py (python)
import flax.nnx as nnx

from utils.nn import STTransformer


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

    def __init__(
        self,
        model_dim: int,
        ffn_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,
        rngs: nnx.Rngs,
    ):
        self.model_dim = model_dim
        self.ffn_dim = ffn_dim
        self.num_latents = num_latents
        self.num_blocks = num_blocks
        self.num_heads = num_heads
        self.dropout = dropout
        self.mask_limit = mask_limit
        self.param_dtype = param_dtype
        self.dtype = dtype
        self.use_flash_attention = use_flash_attention

        self.dynamics = STTransformer(
            self.model_dim,
            self.model_dim,
            self.ffn_dim,
            self.num_latents,
            self.num_blocks,
            self.num_heads,
            self.dropout,
            self.param_dtype,
            self.dtype,
            use_flash_attention=self.use_flash_attention,
            rngs=rngs,
        )
        self.patch_embed = nnx.Embed(self.num_latents, self.model_dim, rngs=rngs)
        self.mask_token = nnx.Param(
            nnx.initializers.lecun_uniform()(rngs.params(), (1, 1, 1, self.model_dim))
        )
        self.action_up = nnx.Linear(
            self.num_latents,
            self.model_dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
            rngs=rngs,
        )

    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.value, vid_embed
            )
[content] models/dynamics.py (python)
            batch_size = vid_embed.shape[0]
            _rng_prob, *_rngs_mask = jax.random.split(batch["mask_rng"], batch_size + 1)
            mask_prob = jax.random.uniform(
                _rng_prob, shape=(batch_size,), minval=self.mask_limit
            )
            per_sample_shape = vid_embed.shape[1:-1]
            mask = jax.vmap(
                lambda rng, prob: jax.random.bernoulli(rng, prob, per_sample_shape),
                in_axes=(0, 0),
            )(jnp.asarray(_rngs_mask), mask_prob)
[tab] models/dynamics.py (python)
from typing import Dict, Any

import jax
import jax.numpy as jnp
import flax.nnx as nnx

from utils.nn import STTransformer


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

    def __init__(
        self,
        model_dim: int,
        ffn_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,
        rngs: nnx.Rngs,
    ):
        self.model_dim = model_dim
        self.ffn_dim = ffn_dim
        self.num_latents = num_latents
        self.num_blocks = num_blocks
        self.num_heads = num_heads
        self.dropout = dropout
        self.mask_limit = mask_limit
        self.param_dtype = param_dtype
        self.dtype = dtype
        self.use_flash_attention = use_flash_attention

        self.dynamics = STTransformer(
            self.model_dim,
            self.model_dim,
            self.ffn_dim,
            self.num_latents,
            self.num_blocks,
            self.num_heads,
            self.dropout,
            self.param_dtype,
            self.dtype,
            use_flash_attention=self.use_flash_attention,
            rngs=rngs,
        )
        self.patch_embed = nnx.Embed(self.num_latents, self.model_dim, rngs=rngs)
        self.mask_token = nnx.Param(
            nnx.initializers.lecun_uniform()(rngs.params(), (1, 1, 1, self.model_dim))
        )
        self.action_up = nnx.Linear(
            self.num_latents,
            self.model_dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
            rngs=rngs,
        )

    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:
            batch_size = vid_embed.shape[0]
            _rng_prob, *_rngs_mask = jax.random.split(batch["mask_rng"], batch_size + 1)
            mask_prob = jax.random.uniform(
                _rng_prob, shape=(batch_size,), minval=self.mask_limit
            )
            per_sample_shape = vid_embed.shape[1:-1]
            mask = jax.vmap(
                lambda rng, prob: jax.random.bernoulli(rng, prob, per_sample_shape),
                in_axes=(0, 0),
            )(jnp.asarray(_rngs_mask), mask_prob)
            mask = mask.at[:, 0].set(False)
            vid_embed = jnp.where(
                jnp.expand_dims(mask, -1), self.mask_token.value, 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)
[tab] experiments/tokenizer_grain_checkpointing.sh (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=8
#SBATCH --time=12:00:00
#SBATCH --cpus-per-task=1
#SBATCH --gres=gpu:8
#SBATCH --mem=100GB

source .venv/bin/activate

data_dir="$PWD/data_arrayrecord/dummy"
ckpt_dir="$PWD/checkpoints/tokenizer_openai_grain_checkpointing_kv_cache"

export XLA_FLAGS=--xla_gpu_autotune_level=0
srun python train_tokenizer.py \
    --batch_size 96 \
    --save_ckpt \
    --log_checkpoint_interval 5 \
    --ckpt_dir $ckpt_dir \
    --num_steps 300000 \
    --warmup_steps 10000 \
    --seed 0 \
    --init_lr=0.0000866 \
    --max_lr=0.0000866 \
    --data_dir $data_dir
[tab] /fast/home/franz.srambical/jafar/utils/parameter_utils.py (python)
from jax.tree_util import tree_map, tree_reduce


def _count_leaf(x):
    """Count parameters in a single leaf node."""
    if hasattr(x, "size"):
        return x.size
    return 0


def _count_component(component_params):
    """Count total parameters in a component."""
    return tree_reduce(
        lambda x, y: x + y, tree_map(_count_leaf, component_params), initializer=0
    )


def count_parameters_by_component(params):
    """Count parameters for each component of the model.

    Args:
        params: Model parameters dictionary
        component_names: List of component names to count. If None, counts all components.

    Returns:
        Dictionary with parameter counts for each component
    """

    component_names = list(params["params"].keys())
    print(f"Counting all components: {component_names}")

    counts = {}
    total_params = 0

    for name in component_names:
        if "params" in params and name in params["params"]:
            component_params = params["params"][name]
        else:
            component_params = params

        count = _count_component(component_params)
        counts[name] = count
        total_params += count

    counts["total"] = total_params
    return counts
[content] /fast/home/franz.srambical/jafar/utils/parameter_utils.py (python)
        params: Model parameters from nnx.split(model, nnx.Param, ...)

    Returns:
        Dictionary with parameter counts for each component
    """
    # In the new NNX API, params is a flat structure with component names as keys
    component_names = list(params.keys())
    print(f"Counting all components: {component_names}")

    counts = {}
    total_params = 0

    for name in component_names:
        component_params = params[name]
[content] utils/nn.py (python)
        attention_mask = jnp.ones((target_seq_len, target_seq_len), dtype=jnp.bool_)
        attention_mask = attention_mask.at[original_seq_len:, :].set(False)
        attention_mask = attention_mask.at[:, original_seq_len:].set(False)

        mask_4d = (
            _fuse_masks(mask, attention_mask) if mask is not None else attention_mask
        )
        mask_4d = mask_4d[jnp.newaxis, jnp.newaxis, :, :]  # (1, 1, seq_len, seq_len)
[content] utils/nn.py (python)
        def _fuse_masks(mask: jax.Array, attention_mask: jax.Array) -> jax.Array:
            mask_bool = mask.astype(jnp.bool_)
            expanded_mask = jnp.pad(
                mask_bool, ((0, pad_size), (0, pad_size)), constant_values=False
            )
            return jnp.logical_and(attention_mask, expanded_mask)
[content] utils/nn.py (python)
        # Filter kwargs to only include valid arguments for dot_product_attention
        valid_kwargs = {
            k: v for k, v in kwargs.items() 
            if k in ['scale', 'query_seq_lengths', 'key_value_seq_lengths', 'local_window_size']
        }
[tab] Environment changes - jafar (markdown)
# Terminal Environment Changes

## Extension: ms-python.debugpy

Enables use of [no-config debugging](https://github.com/microsoft/vscode-python-debugger/wiki/No%E2%80%90Config-Debugging), `debugpy <script.py>`, in the terminal.

- `PYDEVD_DISABLE_FILE_VALIDATION=1`
- `VSCODE_DEBUGPY_ADAPTER_ENDPOINTS=/home/franz.srambical/.cursor-server/extensions/ms-python.debugpy-2025.6.0-linux-x64/.noConfigDebugAdapterEndpoints/endpoint-c77296858146f87c.txt`
- `PATH=${env:PATH}:/home/franz.srambical/.cursor-server/extensions/ms-python.debugpy-2025.6.0-linux-x64/bundled/scripts/noConfigScripts`
- `BUNDLED_DEBUGPY_PATH=/home/franz.srambical/.cursor-server/extensions/ms-python.debugpy-2025.6.0-linux-x64/bundled/libs/debugpy`

## Extension: vscode.git

Enables the following features: git auth provider

- `GIT_ASKPASS=/fast/home/franz.srambical/.cursor-server/bin/a8e95743c5268be73767c46944a71f4465d05c90/extensions/git/dist/askpass.sh`
- `VSCODE_GIT_ASKPASS_NODE=/fast/home/franz.srambical/.cursor-server/bin/a8e95743c5268be73767c46944a71f4465d05c90/node`
- `VSCODE_GIT_ASKPASS_EXTRA_ARGS=`
- `VSCODE_GIT_ASKPASS_MAIN=/fast/home/franz.srambical/.cursor-server/bin/a8e95743c5268be73767c46944a71f4465d05c90/extensions/git/dist/askpass-main.js`
- `VSCODE_GIT_IPC_HANDLE=/run/user/961800067/vscode-git-68e564a388.sock`
[content] train_tokenizer.py (python)
        optimizer_state = nnx.state(optimizer)
        # Create a dummy function to get the abstract shape
        def dummy_fn():
            return optimizer_state
        abstract_optimizer_state = nnx.eval_shape(dummy_fn)
[content] train_tokenizer.py (python)
        optimizer_state = nnx.state(optimizer)
        # Create a function to get the abstract shape
        def get_optimizer_state():
            return optimizer_state
        abstract_optimizer_state = nnx.eval_shape(get_optimizer_state)
[content] /fast/home/franz.srambical/jafar/train_tokenizer.py (python)
        # FIXME (f.srambical): change all checkpointing/ restore logic to separately checkpoint model and optimizer
        optimizer_state = nnx.state(optimizer)
        abstract_optimizer_state = nnx.eval_shape(lambda: optimizer_state)
        restored = checkpoint_manager.restore(
            checkpoint_manager.latest_step(),
            args=ocp.args.Composite(
                model_state=ocp.args.StandardRestore(abstract_optimizer_state),
                # dataloader_state=grain.checkpoint.CheckpointRestore(grain_iterator),
            ),
        )
        restored_optimizer_state = restored["model_state"]
        nnx.update(optimizer, restored_optimizer_state)
        # grain_iterator = restored["dataloader_state"]
[content] /fast/home/franz.srambical/jafar/train_tokenizer.py (python)
        optimizer_state = nnx.state(optimizer)
        abstract_optimizer = nnx.eval_shape(optimizer_state)
        # FIXME (f.srambical): also checkpoint optimizer state
        restored = checkpoint_manager.restore(
            checkpoint_manager.latest_step(),
            args=ocp.args.Composite(
                ocp.args.StandardRestore(abstract_optimizer),
                grain.checkpoint.CheckpointRestore(grain_iterator),
            ),
        )
        restored_optimizer_state = restored[0]
        nnx.update(optimizer, restored_optimizer_state)
        grain_iterator = restored[1]
[content] train_tokenizer.py (python)
        try:
            restored = checkpoint_manager.restore(
                checkpoint_manager.latest_step(),
                args=ocp.args.Composite(
                    ocp.args.StandardRestore(abstract_model_state),
                    grain.checkpoint.CheckpointRestore(grain_iterator),
                ),
            )
            restored_optimizer_state = restored[0]
            nnx.update(optimizer, restored_optimizer_state)
            grain_iterator = restored[1]
        except ValueError as e:
            if "spatial_pos_enc" in str(e) or "temporal_pos_enc" in str(e):
                print("Warning: Checkpoint contains positional encodings that are not in current model.")
                print("This is likely due to a model structure change. Starting from scratch.")
                step = 0
                return
            else:
                raise e
[content] train_tokenizer.py (python)
        restored = checkpoint_manager.restore(
            checkpoint_manager.latest_step(),
            args=ocp.args.Composite(
                model_state=ocp.args.StandardRestore(abstract_model_state),
                dataloader_state=grain.checkpoint.CheckpointRestore(grain_iterator),
            ),
        )
        restored_optimizer_state = restored["model_state"]
        nnx.update(optimizer, restored_optimizer_state)
        grain_iterator = restored["dataloader_state"]
[content] /fast/home/franz.srambical/jafar/train_tokenizer.py (python)
        abstract_model = nnx.eval_shape(optimizer.model)
        abstract_model_state = nnx.state(abstract_model)
        # FIXME (f.srambical): also checkpoint optimizer state
        restored = checkpoint_manager.restore(
            checkpoint_manager.latest_step(),
            args=ocp.args.Composite(
                model_state=ocp.args.StandardRestore(abstract_model_state),
                dataloader_state=grain.checkpoint.CheckpointRestore(grain_iterator),
            ),
        )
        restored_optimizer_state = restored["model_state"]
        nnx.update(optimizer, restored_optimizer_state)
        grain_iterator = restored["dataloader_state"]
        step = checkpoint_manager.latest_step() or 0
        print(f"Restored dataloader and model state from step {step}")
[content] /fast/home/franz.srambical/jafar/train_tokenizer.py (python)
        # Create dummy input for eval_shape
        dummy_videos = jnp.zeros(
            (per_device_batch_size_for_init, args.seq_len, args.image_height, args.image_width, args.image_channels),
            dtype=args.dtype,
        )
        dummy_batch = {"videos": dummy_videos}
        abstract_model = nnx.eval_shape(optimizer.model, dummy_batch)
[content] /fast/home/franz.srambical/jafar/train_tokenizer.py (python)
# Create dummy input for eval_shape
        dummy_videos = jnp.zeros(
            (per_device_batch_size_for_init, args.seq_len, args.image_height, args.image_width, args.image_channels),
            dtype=args.dtype,
        )
        dummy_batch = {"videos": dummy_videos}
[tab] utils/parameter_utils.py (python)
from jax.tree_util import tree_map, tree_reduce


def _count_leaf(x):
    """Count parameters in a single leaf node."""
    if hasattr(x, "size"):
        return x.size
    return 0


def _count_component(component_params):
    """Count total parameters in a component."""
    return tree_reduce(
        lambda x, y: x + y, tree_map(_count_leaf, component_params), initializer=0
    )


def count_parameters_by_component(params):
    """Count parameters for each component of the model.

    Args:
        params: Model parameters from nnx.split(model, nnx.Param, ...)

    Returns:
        Dictionary with parameter counts for each component
    """
    # In the new NNX API, params is a flat structure with component names as keys
    component_names = list(params.keys())
    print(f"Counting all components: {component_names}")

    counts = {}
    total_params = 0

    for name in component_names:
        component_params = params[name]
        count = _count_component(component_params)
        counts[name] = count
        total_params += count

    counts["total"] = total_params
    return counts
[tab] experiments/dynamics_grain_tok_restore.sh (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=2
#SBATCH --ntasks-per-node=8
#SBATCH --time=12:00:00
#SBATCH --cpus-per-task=1
#SBATCH --gres=gpu:8
#SBATCH --mem=100GB

source .venv/bin/activate

data_dir="$PWD/data_arrayrecord/dummy"
ckpt_dir="$PWD/checkpoints/dynamics_openai_grain_tok_lam_restore"
tokenizer_ckpt_dir="$PWD/checkpoints/tokenizer_openai_grain_checkpointing"

export XLA_FLAGS=--xla_gpu_autotune_level=0
srun python train_dynamics.py \
    --save_ckpt \
    --log_checkpoint_interval 5 \
    --batch_size 96 \
    --tokenizer_checkpoint $tokenizer_ckpt_dir \
    --ckpt_dir $ckpt_dir \
    --num_steps 300000 \
    --warmup_steps 10000 \
    --seed 0 \
    --init_lr=0.0000866 \
    --max_lr=0.0000866 \
    --data_dir $data_dir
[tab] slurm/dev/alfred/horeka/overfit_run_ds_oai/lam/from_ckpt/train_lam_dev.sh (shellscript)
# Log the sbatch script
cat $0

module unload mpi/openmpi/5.0
module unload devel/cuda/12.4
source .venv_jafar/bin/activate

WS_DIR='/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/'

TF_RECORDS_DIR=$WS_DIR/data/overfit_dir_openai_tfrecord/test
CHECKPOINT_DIR=$WS_DIR/checkpoints/lam_main_test/${JOB_NAME}_${SLURM_JOB_ID}_redo
mkdir -p $CHECKPOINT_DIR

echo $CHECKPOINT_DIR

env | grep SLURM


batch_size=12
N=1
num_samples=$((batch_size * N))


JOB_NAME=${SLURM_JOB_NAME}_num_samples_${num_samples}
SLURM_JOB_ID=$SLURM_JOB_ID
TAGS="lam overfit_run_oai_ds sample_size_${num_samples} checkpoint_test_baseline main_branch"


echo "num_samples: $num_samples"
echo "job name: $JOB_NAME"

cat train_lam.py

# FROM_CKPT_DIR=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/lam/_3312785/lam_1751470889_161000
# FROM_CKPT_DIR=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/tokenizer/_3316057/lam_1751565277_130

    # --checkpoint=$FROM_CKPT_DIR \

# FROM_CKPT_DIR='/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/tokenizer/_3316657/lam_1751591343_1000'

# FROM_CKPT_DIR='/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/lam/train_lam_overfit_run_oai_ds_samples_12_checkpoint_test_baseline_num_samples_12_3316684/lam_1751593576_156000'
FROM_CKPT_DIR='/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/lam/_3312785/lam_1751470889_193000'

# XLA_FLAGS=--xla_gpu_autotune_level=0 python train_lam.py \
python train_lam.py \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=$batch_size \
    --min_lr=1e-4 \
    --max_lr=1e-4 \
    --log_image_interval=50 \
    --log \
    --name=${JOB_NAME}-${SLURM_JOB_ID}-from_ckpt \
    --tags $TAGS \
    --entity instant-uv \
    --project jafar \
    --data_dir $TF_RECORDS_DIR \
    --num_samples=$num_samples \
    --checkpoint=$FROM_CKPT_DIR \
    --log_checkpoint_interval=500 \
    --model_dim=256 \
    --num_blocks=4 \
    --num_heads=4 \
    --latent_dim=16 \
    --num_latents=6
[tab] /fast/home/franz.srambical/jafar/models/dynamics.py (python)
from typing import Dict, Any

import jax
import jax.numpy as jnp
import flax.nnx as nnx

from utils.nn import STTransformer


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

    def __init__(
        self,
        model_dim: int,
        ffn_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,
        rngs: nnx.Rngs,
    ):
        self.model_dim = model_dim
        self.ffn_dim = ffn_dim
        self.num_latents = num_latents
        self.num_blocks = num_blocks
        self.num_heads = num_heads
        self.dropout = dropout
        self.mask_limit = mask_limit
        self.param_dtype = param_dtype
        self.dtype = dtype
        self.use_flash_attention = use_flash_attention

        self.dynamics = STTransformer(
            self.model_dim,
            self.model_dim,
            self.ffn_dim,
            self.num_latents,
            self.num_blocks,
            self.num_heads,
            self.dropout,
            self.param_dtype,
            self.dtype,
            use_flash_attention=self.use_flash_attention,
            rngs=rngs,
        )
        self.patch_embed = nnx.Embed(self.num_latents, self.model_dim, rngs=rngs)
        self.mask_token = nnx.Param(
            nnx.initializers.lecun_uniform()(rngs.params(), (1, 1, 1, self.model_dim))
        )
        self.action_up = nnx.Linear(
            self.num_latents,
            self.model_dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
            rngs=rngs,
        )

    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:
            batch_size = vid_embed.shape[0]
            _rng_prob, *_rngs_mask = jax.random.split(batch["mask_rng"], batch_size + 1)
            mask_prob = jax.random.uniform(
                _rng_prob, shape=(batch_size,), minval=self.mask_limit
            )
            per_sample_shape = vid_embed.shape[1:-1]
            mask = jax.vmap(
                lambda rng, prob: jax.random.bernoulli(rng, prob, per_sample_shape),
                in_axes=(0, 0),
            )(jnp.asarray(_rngs_mask), mask_prob)
            mask = mask.at[:, 0].set(False)
            vid_embed = jnp.where(
                jnp.expand_dims(mask, -1), self.mask_token.value, 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)
[content] models/lam.py (python)
        # --- Decode ---
        video_recon = self.decoder(video_action_patches)
        video_recon = video_recon.astype(jnp.float32)
        video_recon = nnx.sigmoid(video_recon)
        video_recon = video_recon.astype(self.dtype)
        outputs["recon"] = unpatchify(video_recon, self.patch_size, H, W)
[content] models/lam.py (python)
        self.patch_up = nnx.Linear(
            self.patch_token_dim,
            self.model_dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
            rngs=rngs,
        )
        self.action_up = nnx.Linear(
            self.latent_dim,
            self.model_dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
            rngs=rngs,
        )
        self.decoder = STTransformer(
            self.model_dim,
            self.model_dim,
            self.ffn_dim,
            self.patch_token_dim,
            self.num_blocks,
            self.num_heads,
            self.dropout,
            self.param_dtype,
            self.dtype,
            use_flash_attention=self.use_flash_attention,
            rngs=rngs,
        )
[content] genie.py (python)
        
        # Remove LAM decoder components to save memory (model surgery)
        # Delete decoder components
        del optimizer.model.lam.decoder
        del optimizer.model.lam.patch_up
        del optimizer.model.lam.action_up
        
        # Monkey patch the __call__ method to only do encoding
        def lam_encode_only(self, batch: Dict[str, Any], training: bool = True) -> Dict[str, Any]:
            return self.vq_encode(batch["videos"], training)
        
        optimizer.model.lam.__call__ = lam_encode_only.__get__(optimizer.model.lam, type(optimizer.model.lam))
[tab] experiments/sample.sh (shellscript)
source .venv/bin/activate

data_dir="$PWD/data_arrayrecord/dummy"
ckpt_dir="$PWD/checkpoints/dynamics_openai_grain_tok_lam_restore"

export XLA_FLAGS=--xla_gpu_autotune_level=0
srun python sample.py \
    --batch_size 1 \
    --seq_len 3 \
    --start_frame 1 \
    --checkpoint $ckpt_dir \
    --data_dir $data_dir
[content] sample.py (python)

            checkpoint_manager.latest_step(),
            args=ocp.args.Composite(
                model_state=ocp.args.PyTreeRestore(abstract_optimizer_state),
                dataloader_state=grain.checkpoint.CheckpointRestore(grain_iterator),
            ),
        )
[content] sample.py (python)

    handler_registry = ocp.handlers.DefaultCheckpointHandlerRegistry()
    handler_registry.add(
        "model_state", ocp.args.PyTreeSave, ocp.handlers.PyTreeCheckpointHandler
    )
    handler_registry.add(
        "model_state", ocp.args.PyTreeRestore, ocp.handlers.PyTreeCheckpointHandler
    )
[content] sample.py (python)
# The new Genie.vq_encode returns latent indices directly.
action_batch = genie.vq_encode(batch, training=False)  # type: ignore[arg-type]
action_batch = jnp.asarray(action_batch).reshape(video_batch.shape[0], args.seq_len - 1, 1)
[content] sample.py (python)
# Remove old wrapper and nn.apply usage and replace with nnx.jit sampling function
# --- Define sampling function compatible with new Genie API ---
@nnx.jit
def _sampling_fn(model, batch):
    """Runs Genie.sample with pre-defined generation hyper-parameters."""
    return model.sample(
        batch,
        args.seq_len,
        args.maskgit_steps,
        args.temperature,
        args.sample_argmax,
[content] sample.py (python)
    checkpoint_options = ocp.CheckpointManagerOptions(
        save_interval_steps=args.log_checkpoint_interval,
        max_to_keep=3,
        keep_period=args.log_checkpoint_keep_period,
        step_format_fixed_length=6,
        cleanup_tmp_directories=True,
    )