p-doom/crowd-code-dataset-0.2 huggingface dataset:
[tab] experiments/tokenizer_dummy_distributed_checkpointing_test.sh (shellscript)
#!/usr/bin/env bash
module unload mpi/openmpi/5.0
module unload devel/cuda/12.4
source .venv/bin/activate


data_dir='/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data/dummy'

srun python train_tokenizer.py \
    --batch_size 192 \
    --tmp_ckpt_dir /mnt/odfs/$SLURM_JOB_ID/stripe_default/tokenizer_dummy_distributed_checkpointing_test_tmp \
    --ckpt_dir $PWD/checkpoints/tokenizer_dummy_distributed_checkpointing_test \
    --log_checkpoint_interval 5 \
    --num_steps 300000 \
    --warmup_steps 10000 \
    --seed 0 \
    --min_lr=0.0000866 \
    --max_lr=0.0000866 \
    --data_dir $data_dir
[tab] experiments/dynamics_dummy_test_co_train.sh (shellscript)
module unload mpi/openmpi/5.0
module unload devel/cuda/12.4
source .venv/bin/activate

srun python train_dynamics.py \
    --tokenizer_checkpoint /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/3316022/tokenizer_1751594834_19000 \
    --batch_size=4 \
    --num_steps=50000 \
    --warmup_steps=1250 \
[content] genie.py (python)


    # dummy values since we only use tx to initialize the dummy train states
    dummy_tx = optax.adamw(
        learning_rate=optax.constant_schedule(args.max_lr),
        b1=0.9,
        b2=0.9,
        weight_decay=1e-4,
    )
[content] genie.py (python)


    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)
[content] genie.py (python)
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)
[content] genie.py (python)

    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)
[content] genie.py (python)

    restored_tokenizer_params = (
        PyTreeCheckpointer()
        .restore(
            args.tokenizer_checkpoint,
            item=tokenizer_restore_target,
            restore_args=tokenizer_restore_args,
        )["model"]
        .params["params"]
    )
[tab] experiments/tokenizer_dummy_distributed_checkpointing_test.sh (shellscript)
#!/usr/bin/env bash
module unload mpi/openmpi/5.0
module unload devel/cuda/12.4
source .venv/bin/activate


data_dir='/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data/dummy'

srun python train_tokenizer.py \
    --batch_size 192 \
    --ckpt_dir $PWD/checkpoints/tokenizer_dummy_distributed_checkpointing_test \
    --log_checkpoint_interval 5 \
    --num_steps 300000 \
    --warmup_steps 10000 \
    --seed 0 \
    --min_lr=0.0000866 \
    --max_lr=0.0000866 \
    --data_dir $data_dir
[content] sample.py (python)
gt = video_batch[0, : vid.shape[1]].clip(0, 1).reshape(-1, *video_batch.shape[2:])
recon = vid[0].clip(0, 1).reshape(-1, *vid.shape[2:])
ssim = pix.ssim(gt[args.start_frame + 1 :], recon[args.start_frame + 1 :]).mean()
[content] sample.py (python)
gt = video_batch[:, : vid.shape[1]].clip(0, 1).reshape(-1, *video_batch.shape[2:])
recon = vid.clip(0, 1).reshape(-1, *vid.shape[2:])
ssim = pix.ssim(gt[:, args.start_frame + 1 :], recon[:, args.start_frame + 1 :]).mean()
[content] sample.py (python)
true_video = (video_batch[0] * 255).astype(np.uint8)
pred_video = (vid[0] * 255).astype(np.uint8)
video_comparison = np.zeros((2, *vid.shape[1:5]), dtype=np.uint8)
video_comparison[0] = true_video[: vid.shape[1]]  # Ground truth
video_comparison[1] = pred_video                   # Generated
flat_vid = einops.rearrange(video_comparison, "n t h w c -> t h (n w) c")
[content] genie.py (python)



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
            )
[content] genie.py (python)
    apply_fn=dummy_lam.apply, params=lam_init_params, tx=dummy_tx
        )
        abstract_sharded_lam_state = _create_abstract_sharded_pytree(
            dummy_lam_train_state, sharding
        )
        lam_restore_target = {"model": abstract_sharded_lam_state}
        lam_restore_args = orbax_utils.restore_args_from_target(lam_restore_target)
        restored_lam_params = (
[content] genie.py (python)
    num_heads=args.lam_num_heads,
            dropout=args.dropout,
            codebook_dropout=args.dropout,
        )
        lam_init_params = dummy_lam.init(_rng, inputs)
        dummy_lam_train_state = TrainState.create
[content] genie.py (python)
    # dummy values since we only use tx to initialize the dummy train states
    dummy_tx = optax.adamw(
        learning_rate=optax.constant_schedule(args.max_lr),
        b1=0.9,
        b2=0.9,
        weight_decay=1e-4,
    )
[content] genie.py (python)
dummy_lam = LatentActionModel(
        in_dim=args.image_channels,
        model_dim=args.lam_dim,
        latent_dim=args.latent_patch_dim,
        num_latents=args.num_latent_actions,
        patch_size=args.lam_patch_size,
        num_blocks=args.lam_num_blocks,
        num_heads=args.lam_num_heads,
        dropout=args.dropout,
        codebook_dropout=args.dropout,
    )
[content] genie.py (python)
lam_init_params = dummy_lam.init(_rng, inputs)

    # dummy values since we only use tx to initialize the dummy train states
    dummy_tx = optax.adamw(
        learning_rate=optax.constant_schedule(args.max_lr),
        b1=0.9,
        b2=0.9,
        weight_decay=1e-4,
    )
[content] genie.py (python)
    dummy_lam_train_state = TrainState.create(
        apply_fn=dummy_lam.apply, params=lam_init_params, tx=dummy_tx
    )

    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)
[content] genie.py (python)
            # identify current and future frames to be generated (i.e., t >= step_t)
            future_frame_indicator = (jnp.arange(seq_len) >= step_t) # (S,)
            future_frame_indicator = jnp.broadcast_to(future_frame_indicator[None, :, None], (B, seq_len, N)) # (B, S, N)
            init_mask = future_frame_indicator.astype(bool)
[content] genie.py (python)
            # 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)
            init_mask = mask.astype(bool)
[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

    def setup(self):
        self.dynamics = STTransformer(
            self.model_dim,
            self.num_latents,
            self.num_blocks,
            self.num_heads,
            self.dropout,
        )
        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] genie.py (python)
        # Gather probabilities for the sampled token indices
        final_token_probs = jnp.take_along_axis(
            jax.nn.softmax(final_logits), 
            sampled_token_idxs[..., None], 
            axis=-1
        )[..., 0]
[content] genie.py (python)
        # Get probabilities for the sampled tokens
        final_token_probs = jnp.take_along_axis(
            softmax_probs, 
            sampled_token_idxs[..., None], 
            axis=-1
        )[..., 0]
[content] genie.py (python)
        # Use logits directly for sorting (same order as softmax probabilities)
        # Add large negative values to unmasked positions to ensure they're sorted last
        sort_scores = jnp.where(mask, final_logits, -1e9)
[content] genie.py (python)
        
        # Use jax.lax.cond instead of Python conditional for JAX compatibility
        latent_actions = jax.lax.cond(
            self.lam_co_train,
            lambda: lam_outputs["z_q"],
            lambda: jax.lax.stop_gradient(lam_outputs["z_q"])
        )
[content] genie.py (python)
        rng, _rng = jax.random.split(rng)
        
        def sample_argmax_fn():
            return jnp.argmax(final_logits, axis=-1)
        
        def sample_categorical_fn():
            return jnp.where(
[content] genie.py (python)
        if self.sample_argmax:
            sampled_token_idxs = jnp.argmax(final_logits, axis=-1)
        else:
            rng, _rng = jax.random.split(rng)
            sampled_token_idxs = jnp.where(
[content] genie.py (python)
    # Update tokenizer params using JAX tree operations
    train_state = train_state.replace(
        params=train_state.params.copy(
            params=train_state.params["params"].copy(
                tokenizer=restored_tokenizer_params
            )
        )
    )
[content] genie.py (python)
    def update_tokenizer_params(params_dict):
        if isinstance(params_dict, dict) and "tokenizer" in params_dict:
            return params_dict.copy(tokenizer=restored_tokenizer_params)
        return params_dict
[content] genie.py (python)
    # Update tokenizer params using JAX tree operations
    def update_tokenizer_params(params_dict):
        if isinstance(params_dict, dict) and "tokenizer" in params_dict:
            return params_dict.copy(tokenizer=restored_tokenizer_params)
        return params_dict
    
    train_state = train_state.replace(
        params=jax.tree_util.tree_map(update_tokenizer_params, train_state.params)
    )
[content] genie.py (python)
    # Update tokenizer params by replacing the specific path
    new_params = train_state.params.copy()
    new_params = new_params.copy(
        params=new_params["params"].copy(
            tokenizer=restored_tokenizer_params
        )
    )
    train_state = train_state.replace(params=new_params)
[content] genie.py (python)
        # Update LAM params by replacing the specific path
        new_params = train_state.params.copy()
        new_params = new_params.copy(
            params=new_params["params"].copy(
                lam=restored_lam_params
            )
        )
        train_state = train_state.replace(params=new_params)