p-doom/crowd-code-dataset-0.2 huggingface dataset:
[tab] experiments/train_tokenizer_lr_1e-4.sh (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=2
#SBATCH --ntasks-per-node=4
#SBATCH --time=24:00:00
#SBATCH --partition=accelerated
#SBATCH --cpus-per-task=5
#SBATCH --gres=gpu:4
#SBATCH --output=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/tokenizer-lr-scaling/%x_%j.log
#SBATCH --error=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_/tokenizer-lr-scaling/%x_%j.log
#SBATCH --job-name=train_tokenizer_lr_sweep_1e-4

# Log the sbatch script
cat $0

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

array_records_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data_new/open_ai_minecraft_arrayrecords_chunked

job_name=$SLURM_JOB_NAME
slurm_job_id=$SLURM_JOB_ID

CHECKPOINT_DIR=$PWD/checkpoints/$job_name
mkdir -p $CHECKPOINT_DIR

env | grep SLURM

srun python train_tokenizer.py \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=96 \
    --min_lr=0 \
    --max_lr=1e-4 \
    --log_image_interval=1000 \
    --log_checkpoint_interval=1000 \
    --log \
    --name=mixed-precision-tokenizer-lr-1e-4 \
    --tags tokenizer mixed-precision 1e-4 \
    --entity instant-uv \
    --project jafar \
    --data_dir $array_records_dir
[tab] experiments/train_tokenizer_lr_1e-4_mixed_precision.sh (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=2
#SBATCH --ntasks-per-node=4
#SBATCH --time=24:00:00
#SBATCH --partition=accelerated
#SBATCH --cpus-per-task=5
#SBATCH --gres=gpu:4
#SBATCH --output=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/tokenizer-lr-scaling/%x_%j.log
#SBATCH --error=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_/tokenizer-lr-scaling/%x_%j.log
#SBATCH --job-name=train_tokenizer_lr_sweep_1e-4

# Log the sbatch script
cat $0

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

array_records_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data_new/open_ai_minecraft_arrayrecords_chunked

job_name=$SLURM_JOB_NAME
slurm_job_id=$SLURM_JOB_ID

CHECKPOINT_DIR=$PWD/checkpoints/$job_name
mkdir -p $CHECKPOINT_DIR

env | grep SLURM

srun python train_tokenizer.py \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=96 \
    --min_lr=0 \
    --max_lr=1e-4 \
    --log_image_interval=1000 \
    --log_checkpoint_interval=1000 \
    --log \
    --name=mixed-precision-tokenizer-lr-1e-4 \
    --tags tokenizer mixed-precision 1e-4 \
    --entity instant-uv \
    --project jafar \
    --data_dir $array_records_dir
[content] experiments/train_dynamics_lr_1e-4_mixed_precision.sh (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=2
#SBATCH --ntasks-per-node=4
#SBATCH --time=48:00:00
#SBATCH --partition=accelerated
#SBATCH --cpus-per-task=5
#SBATCH --gres=gpu:4
#SBATCH --output=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_mihir/big-runs/dynamics-cotraining-modelsize-scaling/%x_%j.log
#SBATCH --error=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_mihir/big-runs/dynamics-cotraining-modelsize-scaling/%x_%j.log
#SBATCH --job-name=train_dynamics_modelsize_scaling_36M_2_node

# Log the sbatch script
cat $0

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

array_records_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data_new/open_ai_minecraft_arrayrecords_chunked

job_name=$SLURM_JOB_NAME
slurm_job_id=$SLURM_JOB_ID

CHECKPOINT_DIR=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/dynamics-cotraining-modelsize-scaling/$job_name
mkdir -p $CHECKPOINT_DIR

tokenizer_ckpt_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer-lr-scaling/train_tokenizer_lr_sweep_1e-4

env | grep SLURM

srun python train_dynamics.py \
    --save_ckpt \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=96 \
    --min_lr=0 \
    --max_lr=1.5e-4 \
    --log_image_interval=1000 \
    --log \
    --log_checkpoint_interval=1000 \
    --name=dynamics-debug-run-$slurm_job_id \
    --tags dynamics debug \
    --entity instant-uv \
    --project jafar \
    --tokenizer_checkpoint=$tokenizer_ckpt_dir \
    --data_dir $array_records_dir \
[tab] batch_size_benchmark_tokenizer_mixed_precision.json (json)
{
  "args": {
    "seed": 0,
    "seq_len": 16,
    "image_channels": 3,
    "image_height": 90,
    "image_width": 160,
    "data_dir": "/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data_new/open_ai_minecraft_arrayrecords/10fps_160x90",
    "model_dim": 512,
    "latent_dim": 32,
    "num_latents": 1024,
    "patch_size": 4,
    "num_blocks": 8,
    "num_heads": 8,
    "dropout": 0.0,
    "codebook_dropout": 0.01,
    "dtype":
[content] utils/nn.py (python)
            name='query'
        )(inputs)
        key = nn.Dense(
            self.qkv_features,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
            name='key'
        )(inputs)
        value = nn.Dense(
            self.qkv_features,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
            name='value'
        )(inputs)
        
        # Reshape to (batch_size, seq_len, num_heads, head_dim)
        query = query.reshape(batch_size, seq_len, self.num_heads, head_dim)
        key = key.reshape(batch_size, seq_len, self.num_heads, head_dim)
        value = value.reshape(batch_size, seq_len, self.num_heads, head_dim)
        
        # Transpose to (batch_size, num_heads, seq_len, head_dim) for dot_product_attention
        query = query.transpose(0, 2, 1, 3)
        key = key.transpose(0, 2, 1, 3)
        value = value.transpose(0, 2, 1, 3)
        
        # Prepare mask for dot_product_attention
        attention_mask = None
        if mask is not None:
            # Convert mask to the format expected by dot_product_attention
            if mask.ndim == 2:
                # (seq_len, seq_len) -> (1, 1, seq_len, seq_len)
                attention_mask = mask[None, None, :, :]
            elif mask.ndim == 3:
                # (batch_size, seq_len, seq_len) -> (batch_size, 1, seq_len, seq_len)
                attention_mask = mask[:, None, :, :]
            else:
                # Assume it's already in the right format
                attention_mask = mask
        
        # Apply attention with flash attention if enabled
        implementation = 'cudnn' if self.use_flash_attention else 'xla'
        attention_output = jax.nn.dot_product_attention(
            query=query,
            key=key,
            value=value,
            mask=attention_mask,
            implementation=implementation
        )
        
        # Transpose back to (batch_size, seq_len, num_heads, head_dim)
        attention_output = attention_output.transpose(0, 2, 1, 3)
        
        # Reshape to (batch_size, seq_len, qkv_features)
        attention_output = attention_output.reshape(batch_size, seq_len, self.qkv_features)
        
        # Apply dropout if not deterministic
        if not deterministic and self.dropout_rate > 0.0:
            attention_output = nn.Dropout(
                rate=self.dropout_rate,
                deterministic=deterministic
            )(attention_output)
        
        # Final linear projection
        output = nn.Dense(
            features,
[content] utils/nn.py (python)
    def __call__(self, inputs: jax.Array, mask: Optional[jax.Array] = None, deterministic: bool = False) -> jax.Array:
        """
        Args:
            inputs: Input array of shape (batch_size, seq_len, features)
            mask: Optional mask array of shape (batch_size, seq_len) or (seq_len, seq_len)
            deterministic: Whether to apply dropout
        """
        batch_size, seq_len, features = inputs.shape
        head_dim = self.qkv_features // self.num_heads
        
        # Linear projections for Q, K, V
        query = nn.Dense(
            self.qkv_features,
[content] utils/nn.py (python)


class STBlock(nn.Module):
    dim: int
    num_heads: int
    dropout: float
    param_dtype: jnp.dtype
    dtype: jnp.dtype

    @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 = FlashMultiHeadAttention(
            num_heads=self.num_heads,
            qkv_features=self.dim,
            dropout_rate=self.dropout,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
[content] utils/nn.py (python)


def create_flash_attention_fn(use_flash_attention: bool = True):
    """Create an attention function that uses flash attention if enabled."""
    def attention_fn(query, key, value, bias=None, mask=None, **kwargs):
        implementation = 'cudnn' if use_flash_attention else 'xla'
        return jax.nn.dot_product_attention(
            query=query,
            key=key,
            value=value,
            bias=bias,
            mask=mask,
            implementation=implementation,
            **kwargs
        )
    return attention_fn
[content] utils/nn.py (python)


def _create_flash_attention_fn(use_flash_attention: bool = True, is_causal: bool = False):
    """Create an attention function that uses flash attention if enabled."""
    def attention_fn(query, key, value, bias=None, mask=None, **kwargs):
        implementation = 'cudnn' if use_flash_attention else None
        return jax.nn.dot_product_attention(
            query=query,
            key=key,
            value=value,
            bias=bias,
            mask=mask,
            implementation=implementation,
            is_causal=is_causal,
            **kwargs
        )
    return attention_fn
[content] utils/nn.py (python)

def _create_flash_attention_fn(use_flash_attention: bool, is_causal: bool):
    """Create an attention function that uses flash attention if enabled."""
    def attention_fn(query, key, value, bias=None, mask=None, **kwargs):
        implementation = 'cudnn' if use_flash_attention else None
        return jax.nn.dot_product_attention(
            query=query,
            key=key,
            value=value,
            bias=bias,
            mask=mask,
            implementation=implementation,
            is_causal=is_causal,
            **kwargs
        )
    return attention_fn
[tab] models/tokenizer.py (python)
from typing import Dict, Any, Tuple

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

from utils.preprocess import patchify, unpatchify
from utils.nn import STTransformer, VectorQuantizer


class TokenizerVQVAE(nn.Module):
    """ST-ViVit VQ-VAE"""

    in_dim: int
    model_dim: int
    latent_dim: int
    num_latents: int
    patch_size: int
    num_blocks: int
    num_heads: int
    dropout: float
    codebook_dropout: float
    param_dtype: jnp.dtype
    dtype: jnp.dtype

    def setup(self):
        self.encoder = STTransformer(
            self.model_dim,
            self.latent_dim,
            self.num_blocks,
            self.num_heads,
            self.dropout,
            self.param_dtype,
            self.dtype,
        )
        self.vq = VectorQuantizer(
            self.latent_dim,
            self.num_latents,
            self.codebook_dropout,
        )
        self.out_dim = self.in_dim * self.patch_size**2
        self.decoder = STTransformer(
            self.model_dim,
            self.out_dim,
            self.num_blocks,
            self.num_heads,
            self.dropout,
            self.param_dtype,
            self.dtype,
        )

    def __call__(self, batch: Dict[str, Any], training: bool = True) -> Dict[str, Any]:
        H, W = batch["videos"].shape[2:4]
        outputs = self.vq_encode(batch["videos"], training)
        recon = self.decoder(outputs["z_q"])  # (B, T, H_down * W_down, C)
        recon = recon.astype(jnp.float32)
        recon = nn.sigmoid(recon)
        recon = recon.astype(self.dtype)
        outputs["recon"] = unpatchify(recon, self.patch_size, H, W)
        return outputs

    def vq_encode(self, videos: Any, training: bool = True) -> Dict[str, Any]:
        # --- Preprocess + encode ---
        B, T = videos.shape[:2]
        x = patchify(videos, self.patch_size)
        N = x.shape[2]
        x = self.encoder(x)  # (B, T, N, E)

        # --- Vector quantize ---
        x = x.reshape(B * T * N, self.latent_dim)
        z_q, z, emb, indices = self.vq(x, training)
        z_q = z_q.reshape(B, T, N, self.latent_dim)
        indices = indices.reshape(B, T, N)
        return dict(z_q=z_q, z=z, emb=emb, indices=indices)

    def decode(self, indices: Any, video_hw: Tuple[int, int]):
        z = self.vq.codebook[indices]
        recon = self.decoder(z)
        recon = recon.astype(jnp.float32)
        recon = nn.sigmoid(recon)
        recon = recon.astype(self.dtype)
        return unpatchify(recon, self.patch_size, *video_hw)
[tab] tests/test_checkpointer.py (python)
import unittest
import tempfile
import os
import jax
import jax.numpy as jnp
from flax.training import orbax_utils
from orbax.checkpoint import PyTreeCheckpointer
from pathlib import Path

from models.tokenizer import TokenizerVQVAE
from flax.training.train_state import TrainState
import optax
from jax.sharding import Mesh, PartitionSpec, NamedSharding
from jax.experimental.mesh_utils import create_device_mesh

class DistributedCheckpointerTest(unittest.TestCase):
    def setUp(self):
        super().setUp()
        self._temp_dir_manager = tempfile.TemporaryDirectory()
        self.checkpoint_dir = Path(self._temp_dir_manager.name)
        self.addCleanup(self._temp_dir_manager.cleanup)

        # FIXME (f.srambical): If the tests pass, we should use the default model config instead
        self.model_kwargs = dict(
            in_dim=3,
            model_dim=8,
            latent_dim=4,
            num_latents=16,
            patch_size=2,
            num_blocks=1,
            num_heads=1,
            dropout=0.0,
            codebook_dropout=0.0,
        )
        self.image_shape = (8, 8, 3)
        self.seq_len = 2
        self.batch_size = 2
        self.seed = 0

    def test_distributed_checkpointing(self):
        jax.distributed.initialize()
        num_devices = jax.device_count()
        self.assertGreater(num_devices, 0)

        model = TokenizerVQVAE(**self.model_kwargs)
        rng = jax.random.PRNGKey(self.seed)
        dummy_inputs = dict(
            videos=jnp.zeros((self.batch_size, self.seq_len, *self.image_shape), dtype=jnp.float32)
        )
        params = model.init(rng, dummy_inputs)

        tx = optax.adam(1e-3)
        state = TrainState.create(apply_fn=model.apply, params=params, tx=tx)

        device_mesh_arr = create_device_mesh((num_devices,))
        mesh = Mesh(devices=device_mesh_arr, axis_names=("data",))
        replicated_sharding = NamedSharding(mesh, PartitionSpec())
        state = jax.device_put(state, replicated_sharding)

        ckpt = {"model": state}
        orbax_checkpointer = PyTreeCheckpointer()
        save_args = orbax_utils.save_args_from_target(ckpt)
        ckpt_path = str(self.checkpoint_dir / "test_ckpt")
        orbax_checkpointer.save(ckpt_path, ckpt, save_args=save_args)
        self.assertTrue(os.path.exists(ckpt_path))

        restore_target = {"model": state}
        restore_args = orbax_utils.restore_args_from_target(restore_target)
        restored = orbax_checkpointer.restore(ckpt_path, item=restore_target, restore_args=restore_args)
        for k in state.params.keys():
            self.assertTrue(jax.tree_util.tree_all(jnp.allclose(state.params[k], restored["model"].params[k])))

if __name__ == "__main__":
    unittest.main()
[content] tests/test_checkpointer.py (python)
        self.model_kwargs = {
            "in_dim": 3,
            "model_dim": 8,
            "latent_dim": 4,
            "num_latents": 16,
            "patch_size": 2,
            "num_blocks": 1,
            "num_heads": 1,
            "dropout": 0.0,
            "codebook_dropout": 0.0,
            "param_dtype": jnp.float32,
            "dtype": jnp.bfloat16,
            "use_flash_attention": True,
        }
[tab] models/lam.py (python)
from typing import Dict, Any

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

from utils.preprocess import patchify, unpatchify
from utils.nn import STTransformer, VectorQuantizer


class LatentActionModel(nn.Module):
    """Latent Action ST-ViVit VQ-VAE"""

    in_dim: int
    model_dim: int
    latent_dim: int
    num_latents: int
    patch_size: int
    num_blocks: int
    num_heads: int
    dropout: float
    codebook_dropout: float
    param_dtype: jnp.dtype
    dtype: jnp.dtype

    def setup(self):
        self.patch_token_dim = self.in_dim * self.patch_size**2
        self.encoder = STTransformer(
            self.model_dim,
            self.latent_dim,
            self.num_blocks,
            self.num_heads,
            self.dropout,
            self.param_dtype,
            self.dtype,
        )
        self.action_in = self.param(
            "action_in",
            nn.initializers.lecun_uniform(),
            (1, 1, 1, self.patch_token_dim),
        )
        self.vq = VectorQuantizer(
            self.latent_dim,
            self.num_latents,
            self.codebook_dropout,
        )
        self.patch_up = nn.Dense(
            self.model_dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )
        self.action_up = nn.Dense(
            self.model_dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )
        self.decoder = STTransformer(
            self.model_dim,
            self.patch_token_dim,
            self.num_blocks,
            self.num_heads,
            self.dropout,
            self.param_dtype,
            self.dtype,
        )

    def __call__(self, batch: Dict[str, Any], training: bool = True) -> Dict[str, Any]:
        # --- Encode + VQ ---
        H, W = batch["videos"].shape[2:4]
        outputs = self.vq_encode(batch["videos"], training)
        video_action_patches = self.action_up(outputs["z_q"]) + self.patch_up(
            outputs["patches"][:, :-1]
        )
        del outputs["patches"]

        # --- Decode ---
        video_recon = self.decoder(video_action_patches)
        video_recon = video_recon.astype(jnp.float32)
        video_recon = nn.sigmoid(video_recon)
        video_recon = video_recon.astype(self.dtype)
        outputs["recon"] = unpatchify(video_recon, self.patch_size, H, W)
        return outputs

    def vq_encode(self, videos: Any, training: bool = True) -> Dict[str, Any]:
        # --- Preprocess videos ---
        B, T = videos.shape[:2]
        patches = patchify(videos, self.patch_size)
        action_pad = jnp.broadcast_to(self.action_in, (B, T, 1, self.patch_token_dim))
        padded_patches = jnp.concatenate((action_pad, patches), axis=2)

        # --- Encode ---
        z = self.encoder(padded_patches)  # (B, T, N, E)
        # Get latent action for all future frames
        z = z[:, 1:, 0]  # (B, T-1, E)

        # --- Vector quantize ---
        z = z.reshape(B * (T - 1), self.latent_dim)
        z_q, z, emb, indices = self.vq(z, training)
        z_q = z_q.reshape(B, T - 1, 1, self.latent_dim)
        return dict(patches=patches, z_q=z_q, z=z, emb=emb, indices=indices)
[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
    num_latents: int
    num_blocks: int
    num_heads: int
    dropout: float
    mask_limit: float
    param_dtype: jnp.dtype
    dtype: jnp.dtype

    def setup(self):
        self.dynamics = STTransformer(
            self.model_dim,
            self.num_latents,
            self.num_blocks,
            self.num_heads,
            self.dropout,
            self.param_dtype,
            self.dtype,
        )
        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:
            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)
[tab] tests/test_checkpointer.py (python)
import unittest
import tempfile
import os
import jax
import jax.numpy as jnp
from flax.training import orbax_utils
from orbax.checkpoint import PyTreeCheckpointer
from pathlib import Path

from models.tokenizer import TokenizerVQVAE
from flax.training.train_state import TrainState
import optax
from jax.sharding import Mesh, PartitionSpec, NamedSharding
from jax.experimental.mesh_utils import create_device_mesh

class DistributedCheckpointerTest(unittest.TestCase):
    def setUp(self):
        super().setUp()
        self._temp_dir_manager = tempfile.TemporaryDirectory()
        self.checkpoint_dir = Path(self._temp_dir_manager.name)
        self.addCleanup(self._temp_dir_manager.cleanup)

        # FIXME (f.srambical): If the tests pass, we should use the default model config instead
        self.model_kwargs = {
            "in_dim": 3,
            "model_dim": 8,
            "latent_dim": 4,
            "num_latents": 16,
            "patch_size": 2,
            "num_blocks": 1,
            "num_heads": 1,
            "dropout": 0.0,
            "codebook_dropout": 0.0,
            "param_dtype": jnp.float32,
            "dtype": jnp.bfloat16,
            "use_flash_attention": True,
        }
        self.image_shape = (8, 8, 3)
        self.seq_len = 2
        self.batch_size = 2
        self.seed = 0

    def test_distributed_checkpointing(self):
        jax.distributed.initialize()
        num_devices = jax.device_count()
        self.assertGreater(num_devices, 0)

        model = TokenizerVQVAE(**self.model_kwargs)
        rng = jax.random.PRNGKey(self.seed)
        dummy_inputs = dict(
            videos=jnp.zeros((self.batch_size, self.seq_len, *self.image_shape), dtype=jnp.float32)
        )
        params = model.init(rng, dummy_inputs)

        tx = optax.adam(1e-3)
        state = TrainState.create(apply_fn=model.apply, params=params, tx=tx)

        device_mesh_arr = create_device_mesh((num_devices,))
        mesh = Mesh(devices=device_mesh_arr, axis_names=("data",))
        replicated_sharding = NamedSharding(mesh, PartitionSpec())
        state = jax.device_put(state, replicated_sharding)

        ckpt = {"model": state}
        orbax_checkpointer = PyTreeCheckpointer()
        save_args = orbax_utils.save_args_from_target(ckpt)
        ckpt_path = str(self.checkpoint_dir / "test_ckpt")
        orbax_checkpointer.save(ckpt_path, ckpt, save_args=save_args)
        self.assertTrue(os.path.exists(ckpt_path))

        restore_target = {"model": state}
        restore_args = orbax_utils.restore_args_from_target(restore_target)
        restored = orbax_checkpointer.restore(ckpt_path, item=restore_target, restore_args=restore_args)
        for k in state.params.keys():
            self.assertTrue(jax.tree_util.tree_all(jnp.allclose(state.params[k], restored["model"].params[k])))

if __name__ == "__main__":
    unittest.main()
[tab] models/tokenizer.py (python)
from typing import Dict, Any, Tuple

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

from utils.preprocess import patchify, unpatchify
from utils.nn import STTransformer, VectorQuantizer


class TokenizerVQVAE(nn.Module):
    """ST-ViVit VQ-VAE"""

    in_dim: int
    model_dim: int
    latent_dim: int
    num_latents: int
    patch_size: int
    num_blocks: int
    num_heads: int
    dropout: float
    codebook_dropout: float
    param_dtype: jnp.dtype
    dtype: jnp.dtype
    use_flash_attention: bool

    def setup(self):
        self.encoder = STTransformer(
            self.model_dim,
            self.latent_dim,
            self.num_blocks,
            self.num_heads,
            self.dropout,
            self.param_dtype,
            self.dtype,
            use_flash_attention=self.use_flash_attention,
        )
        self.vq = VectorQuantizer(
            self.latent_dim,
            self.num_latents,
            self.codebook_dropout,
        )
        self.out_dim = self.in_dim * self.patch_size**2
        self.decoder = STTransformer(
            self.model_dim,
            self.out_dim,
            self.num_blocks,
            self.num_heads,
            self.dropout,
            self.param_dtype,
            self.dtype,
            use_flash_attention=self.use_flash_attention,
        )

    def __call__(self, batch: Dict[str, Any], training: bool = True) -> Dict[str, Any]:
        H, W = batch["videos"].shape[2:4]
        outputs = self.vq_encode(batch["videos"], training)
        recon = self.decoder(outputs["z_q"])  # (B, T, H_down * W_down, C)
        recon = recon.astype(jnp.float32)
        recon = nn.sigmoid(recon)
        recon = recon.astype(self.dtype)
        outputs["recon"] = unpatchify(recon, self.patch_size, H, W)
        return outputs

    def vq_encode(self, videos: Any, training: bool = True) -> Dict[str, Any]:
        # --- Preprocess + encode ---
        B, T = videos.shape[:2]
        x = patchify(videos, self.patch_size)
        N = x.shape[2]
        x = self.encoder(x)  # (B, T, N, E)

        # --- Vector quantize ---
        x = x.reshape(B * T * N, self.latent_dim)
        z_q, z, emb, indices = self.vq(x, training)
        z_q = z_q.reshape(B, T, N, self.latent_dim)
        indices = indices.reshape(B, T, N)
        return dict(z_q=z_q, z=z, emb=emb, indices=indices)

    def decode(self, indices: Any, video_hw: Tuple[int, int]):
        z = self.vq.codebook[indices]
        recon = self.decoder(z)
        recon = recon.astype(jnp.float32)
        recon = nn.sigmoid(recon)
        recon = recon.astype(self.dtype)
        return unpatchify(recon, self.patch_size, *video_hw)
[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
    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.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:
            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)
[tab] models/tokenizer.py (python)
from typing import Dict, Any, Tuple

import flax.linen as nn

from utils.preprocess import patchify, unpatchify
from utils.nn import STTransformer, VectorQuantizer


class TokenizerVQVAE(nn.Module):
    """ST-ViVit VQ-VAE"""

    in_dim: int
    model_dim: int
    latent_dim: int
    num_latents: int
    patch_size: int
    num_blocks: int
    num_heads: int
    dropout: float
    codebook_dropout: float
<<<<<<< Updated upstream
=======
    param_dtype: jnp.dtype
    dtype: jnp.dtype
    use_flash_attention: bool
>>>>>>> Stashed changes

    def setup(self):
        self.encoder = STTransformer(
            self.model_dim,
            self.latent_dim,
            self.num_blocks,
            self.num_heads,
            self.dropout,
<<<<<<< Updated upstream
=======
            self.param_dtype,
            self.dtype,
            use_flash_attention=self.use_flash_attention,
>>>>>>> Stashed changes
        )
        self.vq = VectorQuantizer(
            self.latent_dim,
            self.num_latents,
            self.codebook_dropout,
        )
        self.out_dim = self.in_dim * self.patch_size**2
        self.decoder = STTransformer(
            self.model_dim,
            self.out_dim,
            self.num_blocks,
            self.num_heads,
            self.dropout,
<<<<<<< Updated upstream
=======
            self.param_dtype,
            self.dtype,
            use_flash_attention=self.use_flash_attention,
>>>>>>> Stashed changes
        )

    def __call__(self, batch: Dict[str, Any], training: bool = True) -> Dict[str, Any]:
        H, W = batch["videos"].shape[2:4]
        outputs = self.vq_encode(batch["videos"], training)
        recon = self.decoder(outputs["z_q"])  # (B, T, H_down * W_down, C)
        recon = nn.sigmoid(recon)
        outputs["recon"] = unpatchify(recon, self.patch_size, H, W)
        return outputs

    def vq_encode(self, videos: Any, training: bool = True) -> Dict[str, Any]:
        # --- Preprocess + encode ---
        B, T = videos.shape[:2]
        x = patchify(videos, self.patch_size)
        N = x.shape[2]
        x = self.encoder(x)  # (B, T, N, E)

        # --- Vector quantize ---
        x = x.reshape(B * T * N, self.latent_dim)
        z_q, z, emb, indices = self.vq(x, training)
        z_q = z_q.reshape(B, T, N, self.latent_dim)
        indices = indices.reshape(B, T, N)
        return dict(z_q=z_q, z=z, emb=emb, indices=indices)

    def decode(self, indices: Any, video_hw: Tuple[int, int]):
        z = self.vq.codebook[indices]
        recon = self.decoder(z)
        recon = nn.sigmoid(recon)
        return unpatchify(recon, self.patch_size, *video_hw)
[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
    num_latents: int
    num_blocks: int
    num_heads: int
    dropout: float
    mask_limit: float
<<<<<<< Updated upstream
=======
    param_dtype: jnp.dtype
    dtype: jnp.dtype
    use_flash_attention: bool
>>>>>>> Stashed changes

    def setup(self):
        self.dynamics = STTransformer(
            self.model_dim,
            self.num_latents,
            self.num_blocks,
            self.num_heads,
            self.dropout,
<<<<<<< Updated upstream
=======
            self.param_dtype,
            self.dtype,
            use_flash_attention=self.use_flash_attention,
>>>>>>> Stashed changes
        )
        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)

    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)
[content] utils/nn.py (python)
    param_dtype: jnp.dtype
    dtype: jnp.dtype

    @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,
        )(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,
        )(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)
        # FIXME (f.srambical): Here, the attention hidden dimension is the same as the FFN's. Usually, FFN hidden dimension is 4x model_dim
        z = nn.Dense(
            self.dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(z)
        z = nn.gelu(z)
        x = x + z

        return x


class STTransformer(nn.Module):
    model_dim: int
    out_dim: int
    num_blocks: int
    num_heads: int
    dropout: float
    param_dtype: jnp.dtype
    dtype: jnp.dtype

    @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 = STBlock(
                dim=self.model_dim,
                num_heads=self.num_heads,
                dropout=self.dropout,
                param_dtype=self.param_dtype,
                dtype=self.dtype,
            )(x)
        x = nn.Dense(
            self.out_dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )(x)
[content] sample.py (python)


args = tyro.cli(Args)
rng = jax.random.PRNGKey(args.seed)

# --- Load Genie checkpoint ---
genie = Genie(
    # Tokenizer
    in_dim=args.image_channels,
    tokenizer_dim=args.tokenizer_dim,
    latent_patch_dim=args.latent_patch_dim,
    num_patch_latents=args.num_patch_latents,
    patch_size=args.patch_size,
    tokenizer_num_blocks=args.tokenizer_num_blocks,
    tokenizer_num_heads=args.tokenizer_num_heads,
    # LAM
    lam_dim=args.lam_dim,
    latent_action_dim=args.latent_action_dim,
    num_latent_actions=args.num_latent_actions,
    lam_patch_size=args.lam_patch_size,
    lam_num_blocks=args.lam_num_blocks,
    lam_num_heads=args.lam_num_heads,
    # Dynamics
    dyna_dim=args.dyna_dim,
    dyna_num_blocks=args.dyna_num_blocks,
    dyna_num_heads=args.dyna_num_heads,
[content] sample.py (python)
<<<<<<< HEAD
=======
    param_dtype: jnp.dtype = jnp.float32
    dtype: jnp.dtype = jnp.bfloat16
    use_flash_attention: bool = True
>>>>>>> 568891b (feat: use flash attention)


args = tyro.cli(Args)
rng = jax.random.PRNGKey(args.seed)

# --- Load Genie checkpoint ---
genie = Genie(
    # Tokenizer
    in_dim=args.image_channels,
    tokenizer_dim=args.tokenizer_dim,
    latent_patch_dim=args.latent_patch_dim,
    num_patch_latents=args.num_patch_latents,
    patch_size=args.patch_size,
    tokenizer_num_blocks=args.tokenizer_num_blocks,
    tokenizer_num_heads=args.tokenizer_num_heads,
    # LAM
    lam_dim=args.lam_dim,
    latent_action_dim=args.latent_action_dim,
    num_latent_actions=args.num_latent_actions,
    lam_patch_size=args.lam_patch_size,
    lam_num_blocks=args.lam_num_blocks,
    lam_num_heads=args.lam_num_heads,
    # Dynamics
    dyna_dim=args.dyna_dim,
    dyna_num_blocks=args.dyna_num_blocks,
    dyna_num_heads=args.dyna_num_heads,
<<<<<<< HEAD
=======
    param_dtype=args.param_dtype,
    dtype=args.dtype,
    use_flash_attention=args.use_flash_attention,
>>>>>>> 568891b (feat: use flash attention)
[tab] models/tokenizer.py (python)
from typing import Dict, Any, Tuple

import flax.linen as nn

from utils.preprocess import patchify, unpatchify
from utils.nn import STTransformer, VectorQuantizer


class TokenizerVQVAE(nn.Module):
    """ST-ViVit VQ-VAE"""

    in_dim: int
    model_dim: int
    latent_dim: int
    num_latents: int
    patch_size: int
    num_blocks: int
    num_heads: int
    dropout: float
    codebook_dropout: float
<<<<<<< HEAD
=======
    param_dtype: jnp.dtype
    dtype: jnp.dtype
    use_flash_attention: bool
>>>>>>> 568891b (feat: use flash attention)

    def setup(self):
        self.encoder = STTransformer(
            self.model_dim,
            self.latent_dim,
            self.num_blocks,
            self.num_heads,
            self.dropout,
<<<<<<< HEAD
=======
            self.param_dtype,
            self.dtype,
            use_flash_attention=self.use_flash_attention,
>>>>>>> 568891b (feat: use flash attention)
        )
        self.vq = VectorQuantizer(
            self.latent_dim,
            self.num_latents,
            self.codebook_dropout,
        )
        self.out_dim = self.in_dim * self.patch_size**2
        self.decoder = STTransformer(
            self.model_dim,
            self.out_dim,
            self.num_blocks,
            self.num_heads,
            self.dropout,
<<<<<<< HEAD
=======
            self.param_dtype,
            self.dtype,
            use_flash_attention=self.use_flash_attention,
>>>>>>> 568891b (feat: use flash attention)
        )

    def __call__(self, batch: Dict[str, Any], training: bool = True) -> Dict[str, Any]:
        H, W = batch["videos"].shape[2:4]
        outputs = self.vq_encode(batch["videos"], training)
        recon = self.decoder(outputs["z_q"])  # (B, T, H_down * W_down, C)
        recon = nn.sigmoid(recon)
        outputs["recon"] = unpatchify(recon, self.patch_size, H, W)
        return outputs

    def vq_encode(self, videos: Any, training: bool = True) -> Dict[str, Any]:
        # --- Preprocess + encode ---
        B, T = videos.shape[:2]
        x = patchify(videos, self.patch_size)
        N = x.shape[2]
        x = self.encoder(x)  # (B, T, N, E)

        # --- Vector quantize ---
        x = x.reshape(B * T * N, self.latent_dim)
        z_q, z, emb, indices = self.vq(x, training)
        z_q = z_q.reshape(B, T, N, self.latent_dim)
        indices = indices.reshape(B, T, N)
        return dict(z_q=z_q, z=z, emb=emb, indices=indices)

    def decode(self, indices: Any, video_hw: Tuple[int, int]):
        z = self.vq.codebook[indices]
        recon = self.decoder(z)
        recon = nn.sigmoid(recon)
        return unpatchify(recon, self.patch_size, *video_hw)
[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
    num_latents: int
    num_blocks: int
    num_heads: int
    dropout: float
    mask_limit: float
<<<<<<< HEAD
=======
    param_dtype: jnp.dtype
    dtype: jnp.dtype
    use_flash_attention: bool
>>>>>>> 568891b (feat: use flash attention)

    def setup(self):
        self.dynamics = STTransformer(
            self.model_dim,
            self.num_latents,
            self.num_blocks,
            self.num_heads,
            self.dropout,
<<<<<<< HEAD
=======
            self.param_dtype,
            self.dtype,
            use_flash_attention=self.use_flash_attention,
>>>>>>> 568891b (feat: 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)

    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)
[tab] models/tokenizer.py (python)
from typing import Dict, Any, Tuple

import flax.linen as nn

from utils.preprocess import patchify, unpatchify
from utils.nn import STTransformer, VectorQuantizer


class TokenizerVQVAE(nn.Module):
    """ST-ViVit VQ-VAE"""

    in_dim: int
    model_dim: int
    latent_dim: int
    num_latents: int
    patch_size: int
    num_blocks: int
    num_heads: int
    dropout: float
    codebook_dropout: float
    use_flash_attention: bool

    def setup(self):
        self.encoder = STTransformer(
            self.model_dim,
            self.latent_dim,
            self.num_blocks,
            self.num_heads,
            self.dropout,
            use_flash_attention=self.use_flash_attention,
        )
        self.vq = VectorQuantizer(
            self.latent_dim,
            self.num_latents,
            self.codebook_dropout,
        )
        self.out_dim = self.in_dim * self.patch_size**2
        self.decoder = STTransformer(
            self.model_dim,
            self.out_dim,
            self.num_blocks,
            self.num_heads,
            self.dropout,
            use_flash_attention=self.use_flash_attention,
        )

    def __call__(self, batch: Dict[str, Any], training: bool = True) -> Dict[str, Any]:
        H, W = batch["videos"].shape[2:4]
        outputs = self.vq_encode(batch["videos"], training)
        recon = self.decoder(outputs["z_q"])  # (B, T, H_down * W_down, C)
        recon = nn.sigmoid(recon)
        outputs["recon"] = unpatchify(recon, self.patch_size, H, W)
        return outputs

    def vq_encode(self, videos: Any, training: bool = True) -> Dict[str, Any]:
        # --- Preprocess + encode ---
        B, T = videos.shape[:2]
        x = patchify(videos, self.patch_size)
        N = x.shape[2]
        x = self.encoder(x)  # (B, T, N, E)

        # --- Vector quantize ---
        x = x.reshape(B * T * N, self.latent_dim)
        z_q, z, emb, indices = self.vq(x, training)
        z_q = z_q.reshape(B, T, N, self.latent_dim)
        indices = indices.reshape(B, T, N)
        return dict(z_q=z_q, z=z, emb=emb, indices=indices)

    def decode(self, indices: Any, video_hw: Tuple[int, int]):
        z = self.vq.codebook[indices]
        recon = self.decoder(z)
        recon = nn.sigmoid(recon)
        return unpatchify(recon, self.patch_size, *video_hw)