p-doom/crowd-code-dataset-0.2 huggingface dataset:
[content] utils/nn.py (python)

    """
    Dimension keys:
        B: batch size
        T: number of frames
        N: number of patches per frame
        I: number of input features
        M: model dimension
        D: FFN dimension
        O: number of output features
        F: number of frames
        P: number of patches
    """
[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/maskgit_dynamics_openai_grain_tok_restore"
tokenizer_ckpt_dir="$PWD/checkpoints/tokenizer_openai_grain_checkpointing"

export XLA_FLAGS=--xla_gpu_autotune_level=0
export PYTHONUNBUFFERED=1
srun python train_dynamics.py \
    --save_ckpt \
    --log_checkpoint_interval 5 \
    --batch_size 12 \
    --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
[content] genie.py (python)

    def sample(self, batch: Dict[str, jax.Array]) -> jax.Array:
         # FIXME (f.srambical): implement this
        raise NotImplementedError("Sampling not implemented for fully causal dynamics model")
[content] genie.py (python)
    def vq_encode(self, batch: Dict[str, jax.Array], training: bool) -> jax.Array:
        # --- Preprocess videos ---
        video_BTHWC = batch["videos"]
        lam_output = self.lam.vq_encode(video_BTHWC, training=training)
        lam_indices_E = lam_output["indices"]
        return lam_indices_E
[content] genie.py (python)

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