p-doom/crowd-code-dataset-0.2 huggingface dataset:
[tab] profile_dynamics.patch (diff)
diff --git a/train_dynamics.py b/train_dynamics.py
index cd9f431..ae4af40 100644
--- a/train_dynamics.py
+++ b/train_dynamics.py
@@ -134,7 +134,7 @@ def dynamics_loss_fn(
     return ce_loss, (outputs["recon"], metrics)
 
 
-@nnx.jit
+# @nnx.jit
 def train_step(
     model: Genie, optimizer: nnx.Optimizer, inputs: dict
 ) -> tuple[jax.Array, jax.Array, dict]:
@@ -353,16 +353,15 @@ if __name__ == "__main__":
         del optimizer.model.tokenizer.vq.drop
 
     # --- TRAIN LOOP ---
-    dataloader = (
-        jax.make_array_from_process_local_data(videos_sharding, elem)
-        for elem in grain_iterator
-    )
     print(f"Starting training from step {step}...")
-    while step < args.num_steps:
-        for videos in dataloader:
+    should_break = False
+    while step < args.num_steps and not should_break:
+        for _ in range(5):
+            if step == 2:
+                jax.profiler.start_trace("tensorboard")
             # --- Train step ---
             rng, _rng_mask = jax.random.split(rng, 2)
-            inputs = dict(videos=videos, mask_rng=_rng_mask)
+            inputs = dict(mask_rng=_rng_mask)
             loss, recon, metrics = train_step(optimizer.model, optimizer, inputs)
             metrics["lr"] = lr_schedule(step)
             print(f"Step {step}, loss: {loss}")
@@ -379,21 +378,7 @@ if __name__ == "__main__":
                         }
                     )
                 if step % args.log_image_interval == 0:
-                    gt_seq = inputs["videos"][0].astype(jnp.float32) / 255.0
-                    recon_seq = recon[0].clip(0, 1)
-                    comparison_seq = jnp.concatenate((gt_seq, recon_seq), axis=1)
-                    comparison_seq = einops.rearrange(
-                        comparison_seq * 255, "t h w c -> h (t w) c"
-                    )
-                    if jax.process_index() == 0:
-                        log_images = dict(
-                            image=wandb.Image(np.asarray(gt_seq[args.seq_len - 1])),
-                            recon=wandb.Image(np.asarray(recon_seq[args.seq_len - 1])),
-                            true_vs_recon=wandb.Image(
-                                np.asarray(comparison_seq.astype(np.uint8))
-                            ),
-                        )
-                        wandb.log(log_images)
+                    pass
             # --- Checkpointing ---
             if args.save_ckpt and step % args.log_checkpoint_interval == 0:
                 optimizer_state = nnx.state(optimizer)
[content] train_dynamics.py (python)
) -> tuple[jax.Array, tuple[jax.Array | None, dict]]:
    """Compute masked dynamics loss"""
    model.train()
    outputs = model(inputs, training=True)
    mask = outputs["mask"]
    outputs["token_logits"] = outputs["token_logits"].astype(jnp.float32)
    ce_loss = optax.softmax_cross_entropy_with_integer_labels(
        outputs["token_logits"], outputs["video_tokens"]
    )
    ce_loss = (mask * ce_loss).sum() / mask.sum()
    acc = outputs["token_logits"].argmax(-1) == outputs["video_tokens"]
    acc = (mask * acc).sum() / mask.sum()
    select_probs = jax.nn.softmax(outputs["token_logits"])
    _, index_counts_tokenizer = jnp.unique_counts(
        jnp.ravel(outputs["video_tokens"]), size=args.num_patch_latents, fill_value=0
    )
    codebook_usage_tokenizer = (index_counts_tokenizer != 0).mean()
    metrics = dict(
        cross_entropy_loss=ce_loss,
        masked_token_accuracy=acc,
        select_logit=outputs["token_logits"].max(-1).mean(),
        select_p=select_probs.max(-1).mean(),
        entropy=jax.scipy.special.entr(select_probs).sum(-1).mean(),
        codebook_usage_tokenizer=codebook_usage_tokenizer,
    )
    return ce_loss, (None, metrics)


@nnx.jit
def train_step(
    model: Genie, optimizer: nnx.Optimizer, inputs: dict
) -> tuple[jax.Array, jax.Array | None, dict]:
    """Update state and compute metrics"""

    def loss_fn(model: Genie) -> tuple[jax.Array, tuple[jax.Array | None, dict]]:
[content] train_lam.py (python)
    # Use model configuration where possible; for vq beta, read from model if present,
    # otherwise default to 0.25 (same as Args default)
    vq_beta = getattr(model, "vq_beta", 0.25)
    loss = mse + q_loss + vq_beta * commitment_loss
[content] train_tokenizer.py (python)
        (loss, (recon, metrics)), grads = nnx.value_and_grad(loss_fn, has_aux=True)(
            tokenizer
        )
        optimizer.update(grads)
        if args.log_gradients:
            metrics["encoder_gradients_std/"] = jax.tree.map(
                lambda x: x.std(), grads["params"]["encoder"]
            )
            metrics["vq_gradients_std/"] = jax.tree.map(
                lambda x: x.std(), grads["params"]["vq"]
            )
            metrics["decoder_gradients_std/"] = jax.tree.map(
                lambda x: x.std(), grads["params"]["decoder"]
            )
        return loss, recon, metrics
[content] train_tokenizer.py (python)
    @nnx.jit
    def train_step(
        tokenizer: TokenizerVQVAE, optimizer: nnx.Optimizer, inputs: dict
    ) -> tuple[jax.Array, jax.Array, dict]:
        def loss_fn(model: TokenizerVQVAE) -> tuple[jax.Array, tuple[jax.Array, dict]]:
            return tokenizer_loss_fn(model, inputs)
[content] train_tokenizer.py (python)
        # --- Compute validation metrics ---
        gt_clipped = gt.clip(0, 1).reshape(-1, *gt.shape[2:])
        recon = outputs["recon"].clip(0, 1).reshape(-1, *outputs["recon"].shape[2:])
        psnr = jnp.asarray(pix.psnr(gt_clipped, recon)).mean()
        ssim = jnp.asarray(pix.ssim(gt_clipped, recon)).mean()
        _, index_counts = jnp.unique_counts(
            jnp.ravel(outputs["indices"]), size=args.num_latents, fill_value=0
        )
        codebook_usage = (index_counts != 0).mean()
        metrics = dict(
            loss=loss,
            mse=mse,
            q_loss=q_loss,
            commitment_loss=commitment_loss,
            psnr=psnr,
            ssim=ssim,
            codebook_usage=codebook_usage,
        )
        return loss, (outputs["recon"], metrics)
[content] train_tokenizer.py (python)
def make_tokenizer_loss_fn(args: Args):
    def tokenizer_loss_fn(
        model: TokenizerVQVAE, inputs: dict
    ) -> tuple[jax.Array, tuple[jax.Array, dict]]:
        # --- Compute loss ---
        # FIXME (f.srambical): Can we even do native int8 training without casting the video at all?
        # FIXME (f.srambical): If the tokenizer is the reason for the dynamics model being memory-bound,
        # should we at least train the tokenizer natively in int8?
        gt = jnp.asarray(inputs["videos"], dtype=jnp.float32) / 255.0
        inputs["videos"] = gt.astype(args.dtype)
        model.train()
        outputs = model(inputs, training=True)
        outputs["recon"] = outputs["recon"].astype(jnp.float32)
        mse = jnp.square(gt - outputs["recon"]).mean()
        q_loss = jnp.square(jax.lax.stop_gradient(outputs["emb"]) - outputs["z"]).mean()
        commitment_loss = jnp.square(
            outputs["emb"] - jax.lax.stop_gradient(outputs["z"])
        ).mean()
        loss = mse + q_loss + args.vq_beta * commitment_loss
[content] train_tokenizer.py (python)
    (loss, (recon, metrics)), grads = nnx.value_and_grad(loss_fn, has_aux=True)(
        tokenizer
    )
    optimizer.update(grads)
    if args.log_gradients:
        metrics["encoder_gradients_std/"] = jax.tree.map(
            lambda x: x.std(), grads["params"]["encoder"]
        )
        metrics["vq_gradients_std/"] = jax.tree.map(
            lambda x: x.std(), grads["params"]["vq"]
        )
        metrics["decoder_gradients_std/"] = jax.tree.map(
            lambda x: x.std(), grads["params"]["decoder"]
        )
    return loss, recon, metrics
[content] train_tokenizer.py (python)
@nnx.jit
def train_step(
    tokenizer: TokenizerVQVAE, optimizer: nnx.Optimizer, inputs: dict, args: Args
) -> tuple[jax.Array, jax.Array, dict]:
    def loss_fn(model: TokenizerVQVAE) -> tuple[jax.Array, tuple[jax.Array, dict]]:
        return tokenizer_loss_fn(model, inputs, args)
[content] train_tokenizer.py (python)
    # --- Compute validation metrics ---
    gt = gt.clip(0, 1).reshape(-1, *gt.shape[2:])
    recon = outputs["recon"].clip(0, 1).reshape(-1, *outputs["recon"].shape[2:])
    psnr = jnp.asarray(pix.psnr(gt, recon)).mean()
    ssim = jnp.asarray(pix.ssim(gt, recon)).mean()
    _, index_counts = jnp.unique_counts(
        jnp.ravel(outputs["indices"]), size=args.num_latents, fill_value=0
    )
    codebook_usage = (index_counts != 0).mean()
    metrics = dict(
        loss=loss,
        mse=mse,
        q_loss=q_loss,
        commitment_loss=commitment_loss,
        psnr=psnr,
        ssim=ssim,
        codebook_usage=codebook_usage,
    )
    return loss, (outputs["recon"], metrics)
[content] train_tokenizer.py (python)
def tokenizer_loss_fn(
    model: TokenizerVQVAE, inputs: dict, args: Args
) -> tuple[jax.Array, tuple[jax.Array, dict]]:
    # --- Compute loss ---
    # FIXME (f.srambical): Can we even do native int8 training without casting the video at all?
    # FIXME (f.srambical): If the tokenizer is the reason for the dynamics model being memory-bound,
    # should we at least train the tokenizer natively in int8?
    gt = jnp.asarray(inputs["videos"], dtype=jnp.float32) / 255.0
    inputs["videos"] = gt.astype(args.dtype)
    model.train()
    outputs = model(inputs, training=True)
    outputs["recon"] = outputs["recon"].astype(jnp.float32)
    mse = jnp.square(gt - outputs["recon"]).mean()
    q_loss = jnp.square(jax.lax.stop_gradient(outputs["emb"]) - outputs["z"]).mean()
    commitment_loss = jnp.square(
        outputs["emb"] - jax.lax.stop_gradient(outputs["z"])
    ).mean()
    loss = mse + q_loss + args.vq_beta * commitment_loss
[content] train_tokenizer.py (python)
def tokenizer_loss_fn(
    model: TokenizerVQVAE, inputs: dict
) -> tuple[jax.Array, tuple[jax.Array, dict]]:
    # --- Compute loss ---
    # FIXME (f.srambical): Can we even do native int8 training without casting the video at all?
    # FIXME (f.srambical): If the tokenizer is the reason for the dynamics model being memory-bound,
    # should we at least train the tokenizer natively in int8?
    gt = jnp.asarray(inputs["videos"], dtype=jnp.float32) / 255.0
    inputs["videos"] = gt.astype(args.dtype)
    model.train()
    outputs = model(inputs, training=True)
    outputs["recon"] = outputs["recon"].astype(jnp.float32)
    mse = jnp.square(gt - outputs["recon"]).mean()
    q_loss = jnp.square(jax.lax.stop_gradient(outputs["emb"]) - outputs["z"]).mean()
    commitment_loss = jnp.square(
        outputs["emb"] - jax.lax.stop_gradient(outputs["z"])
    ).mean()
    loss = mse + q_loss + args.vq_beta * commitment_loss

    # --- Compute validation metrics ---
    gt = gt.clip(0, 1).reshape(-1, *gt.shape[2:])
    recon = outputs["recon"].clip(0, 1).reshape(-1, *outputs["recon"].shape[2:])
    psnr = jnp.asarray(pix.psnr(gt, recon)).mean()
    ssim = jnp.asarray(pix.ssim(gt, recon)).mean()
    _, index_counts = jnp.unique_counts(
        jnp.ravel(outputs["indices"]), size=args.num_latents, fill_value=0
    )
    codebook_usage = (index_counts != 0).mean()
    metrics = dict(
        loss=loss,
        mse=mse,
        q_loss=q_loss,
        commitment_loss=commitment_loss,
        psnr=psnr,
        ssim=ssim,
        codebook_usage=codebook_usage,
    )
    return loss, (outputs["recon"], metrics)


@nnx.jit
def train_step(
    tokenizer: TokenizerVQVAE, optimizer: nnx.Optimizer, inputs: dict
) -> tuple[jax.Array, jax.Array, dict]:
    def loss_fn(model: TokenizerVQVAE) -> tuple[jax.Array, tuple[jax.Array, dict]]:
        return tokenizer_loss_fn(model, inputs)

    (loss, (recon, metrics)), grads = nnx.value_and_grad(loss_fn, has_aux=True)(
        tokenizer
    )
    optimizer.update(grads)
    if args.log_gradients:
        metrics["encoder_gradients_std/"] = jax.tree.map(
            lambda x: x.std(), grads["params"]["encoder"]
        )
        metrics["vq_gradients_std/"] = jax.tree.map(
            lambda x: x.std(), grads["params"]["vq"]
        )
        metrics["decoder_gradients_std/"] = jax.tree.map(
            lambda x: x.std(), grads["params"]["decoder"]
        )
    return loss, recon, metrics
[content] train_tokenizer.py (python)
    # --- Define loss and train step (close over args) ---
    def tokenizer_loss_fn(
        model: TokenizerVQVAE, inputs: dict
    ) -> tuple[jax.Array, tuple[jax.Array, dict]]:
        gt = jnp.asarray(inputs["videos"], dtype=jnp.float32) / 255.0
        inputs["videos"] = gt.astype(args.dtype)
        model.train()
        outputs = model(inputs, training=True)
        outputs["recon"] = outputs["recon"].astype(jnp.float32)
        mse = jnp.square(gt - outputs["recon"]).mean()
        q_loss = jnp.square(
            jax.lax.stop_gradient(outputs["emb"]) - outputs["z"]
        ).mean()
        commitment_loss = jnp.square(
            outputs["emb"] - jax.lax.stop_gradient(outputs["z"])
        ).mean()
        loss = mse + q_loss + args.vq_beta * commitment_loss

        gt_clipped = gt.clip(0, 1).reshape(-1, *gt.shape[2:])
        recon = outputs["recon"].clip(0, 1).reshape(-1, *outputs["recon"].shape[2:])
        psnr = jnp.asarray(pix.psnr(gt_clipped, recon)).mean()
        ssim = jnp.asarray(pix.ssim(gt_clipped, recon)).mean()
        _, index_counts = jnp.unique_counts(
            jnp.ravel(outputs["indices"]), size=args.num_latents, fill_value=0
        )
        codebook_usage = (index_counts != 0).mean()
        metrics = dict(
            loss=loss,
            mse=mse,
            q_loss=q_loss,
            commitment_loss=commitment_loss,
            psnr=psnr,
            ssim=ssim,
            codebook_usage=codebook_usage,
        )
        return loss, (outputs["recon"], metrics)

    @nnx.jit
    def train_step(
        tokenizer: TokenizerVQVAE, optimizer: nnx.Optimizer, inputs: dict
    ) -> tuple[jax.Array, jax.Array, dict]:
        def loss_fn(model: TokenizerVQVAE) -> tuple[jax.Array, tuple[jax.Array, dict]]:
            return tokenizer_loss_fn(model, inputs)

        (loss, (recon, metrics)), grads = nnx.value_and_grad(loss_fn, has_aux=True)(
            tokenizer
        )
        optimizer.update(grads)
        if args.log_gradients:
            metrics["encoder_gradients_std/"] = jax.tree.map(
                lambda x: x.std(), grads["params"]["encoder"]
            )
            metrics["vq_gradients_std/"] = jax.tree.map(
                lambda x: x.std(), grads["params"]["vq"]
            )
            metrics["decoder_gradients_std/"] = jax.tree.map(
                lambda x: x.std(), grads["params"]["decoder"]
            )
        return loss, recon, metrics
[content] train_lam.py (python)
def lam_loss_fn(
    model: LatentActionModel, inputs: dict, args: Args
) -> tuple[jax.Array, tuple[jax.Array, jax.Array, dict]]:
    # --- Compute loss ---
    gt = jnp.asarray(inputs["videos"], dtype=jnp.float32) / 255.0
    inputs["videos"] = gt.astype(args.dtype)
    model.train()
    outputs = model(inputs, training=True)
    outputs["recon"] = outputs["recon"].astype(jnp.float32)
    gt_future_frames = gt[:, 1:]
    mse = jnp.square(gt_future_frames - outputs["recon"]).mean()
    q_loss = jnp.square(jax.lax.stop_gradient(outputs["emb"]) - outputs["z"]).mean()
    commitment_loss = jnp.square(
        outputs["emb"] - jax.lax.stop_gradient(outputs["z"])
    ).mean()
    loss = mse + q_loss + args.vq_beta * commitment_loss

    # --- Compute validation metrics ---
    gt = gt_future_frames.clip(0, 1).reshape(-1, *gt_future_frames.shape[2:])
    recon = outputs["recon"].clip(0, 1).reshape(-1, *outputs["recon"].shape[2:])
    psnr = jnp.asarray(pix.psnr(gt, recon)).mean()
    ssim = jnp.asarray(pix.ssim(gt, recon)).mean()
    count_fn = jax.vmap(lambda i: (outputs["indices"] == i).sum())
    index_counts = count_fn(jnp.arange(args.num_latents))
    metrics = dict(
        loss=loss,
        mse=mse,
        q_loss=q_loss,
        commitment_loss=commitment_loss,
        psnr=psnr,
        ssim=ssim,
        codebook_usage=(index_counts != 0).mean(),
    )
    return loss, (outputs["recon"], index_counts, metrics)


@nnx.jit
def train_step(
    lam: LatentActionModel,
    optimizer: nnx.Optimizer,
    inputs: dict,
    action_last_active: jax.Array,
    rng: jax.Array,
    args: Args,
) -> tuple[jax.Array, jax.Array, jax.Array, dict]:
    def loss_fn(
        model: LatentActionModel,
    ) -> tuple[jax.Array, tuple[jax.Array, jax.Array, dict]]:
        return lam_loss_fn(model, inputs, args)

    # --- Update model ---
    (loss, (recon, idx_counts, metrics)), grads = nnx.value_and_grad(
        loss_fn, has_aux=True
    )(lam)
    optimizer.update(grads)

    # --- Reset inactive latent actions ---
    codebook = lam.vq.codebook
    num_codes = len(codebook)
    active_codes = idx_counts != 0.0
    action_last_active = jnp.where(active_codes, 0, action_last_active + 1)
    p_code = active_codes / active_codes.sum()
    reset_idxs = jax.random.choice(rng, num_codes, shape=(num_codes,), p=p_code)
    do_reset = action_last_active >= args.vq_reset_thresh
    new_codebook = jnp.where(
        jnp.expand_dims(do_reset, -1), codebook[reset_idxs], codebook.value
    )
    lam.vq.codebook.value = new_codebook
    action_last_active = jnp.where(do_reset, 0, action_last_active)
    return loss, recon, action_last_active, metrics
[content] train_lam.py (python)
    # --- Define loss and train step (close over args) ---
    def lam_loss_fn(
        model: LatentActionModel, inputs: dict
    ) -> tuple[jax.Array, tuple[jax.Array, jax.Array, dict]]:
        gt = jnp.asarray(inputs["videos"], dtype=jnp.float32) / 255.0
        inputs["videos"] = gt.astype(args.dtype)
        model.train()
        outputs = model(inputs, training=True)
        outputs["recon"] = outputs["recon"].astype(jnp.float32)
        gt_future_frames = gt[:, 1:]
        mse = jnp.square(gt_future_frames - outputs["recon"]).mean()
        q_loss = jnp.square(
            jax.lax.stop_gradient(outputs["emb"]) - outputs["z"]
        ).mean()
        commitment_loss = jnp.square(
            outputs["emb"] - jax.lax.stop_gradient(outputs["z"])
        ).mean()
        loss = mse + q_loss + args.vq_beta * commitment_loss

        gt_val = gt_future_frames.clip(0, 1).reshape(-1, *gt_future_frames.shape[2:])
        recon = outputs["recon"].clip(0, 1).reshape(-1, *outputs["recon"].shape[2:])
        psnr = jnp.asarray(pix.psnr(gt_val, recon)).mean()
        ssim = jnp.asarray(pix.ssim(gt_val, recon)).mean()
        count_fn = jax.vmap(lambda i: (outputs["indices"] == i).sum())
        index_counts = count_fn(jnp.arange(args.num_latents))
        metrics = dict(
            loss=loss,
            mse=mse,
            q_loss=q_loss,
            commitment_loss=commitment_loss,
            psnr=psnr,
            ssim=ssim,
            codebook_usage=(index_counts != 0).mean(),
        )
        return loss, (outputs["recon"], index_counts, metrics)

    @nnx.jit
    def train_step(
        lam: LatentActionModel,
        optimizer: nnx.Optimizer,
        inputs: dict,
        action_last_active: jax.Array,
        rng: jax.Array,
    ) -> tuple[jax.Array, jax.Array, jax.Array, dict]:
        def loss_fn(
            model: LatentActionModel,
        ) -> tuple[jax.Array, tuple[jax.Array, jax.Array, dict]]:
            return lam_loss_fn(model, inputs)

        # --- Update model ---
        (loss, (recon, idx_counts, metrics)), grads = nnx.value_and_grad(
            loss_fn, has_aux=True
        )(lam)
        optimizer.update(grads)

        # --- Reset inactive latent actions ---
        codebook = lam.vq.codebook
        num_codes = len(codebook)
        active_codes = idx_counts != 0.0
        action_last_active = jnp.where(active_codes, 0, action_last_active + 1)
        p_code = active_codes / active_codes.sum()
        reset_idxs = jax.random.choice(rng, num_codes, shape=(num_codes,), p=p_code)
        do_reset = action_last_active >= args.vq_reset_thresh
        new_codebook = jnp.where(
            jnp.expand_dims(do_reset, -1), codebook[reset_idxs], codebook.value
        )
        lam.vq.codebook.value = new_codebook
        action_last_active = jnp.where(do_reset, 0, action_last_active)
        return loss, recon, action_last_active, metrics
[content] train_dynamics.py (python)
    jax.config.update("jax_transfer_guard", "disallow")
    should_break = False
    while step < args.num_steps and not should_break:
        for _ in range(5):
            # if step == 2:
            #     jax.profiler.start_trace("tensorboard")
[content] train_dynamics.py (python)
    # --- Create DataLoaderIterator from dataloader ---
    image_shape = (args.image_height, args.image_width, args.image_channels)
    array_record_files = [
        os.path.join(args.data_dir, x)
        for x in os.listdir(args.data_dir)
        if x.endswith(".array_record")
    ]
    grain_dataloader = get_dataloader(
        array_record_files,
        args.seq_len,
        # NOTE: We deliberately pass the global batch size
        # The dataloader shards the dataset across all processes
        args.batch_size,
        *image_shape,
        num_workers=8,
        prefetch_buffer_size=1,
        seed=args.seed,
    )
    initial_state = grain_dataloader._create_initial_state()
    grain_iterator = grain.DataLoaderIterator(grain_dataloader, initial_state)
[content] train_dynamics.py (python)
    if args.log and jax.process_index() == 0:
        wandb_init_kwargs = {
            "entity": args.entity,
            "project": args.project,
            "name": args.name,
            "tags": args.tags,
            "group": "debug",
            "config": args,
        }

        if args.wandb_id:
            wandb_init_kwargs.update(
                {
                    "id": args.wandb_id,
                    "resume": "allow",
                }
            )
        wandb.init(**wandb_init_kwargs)

        wandb.config.update({"model_param_count": param_counts})

    print("Parameter counts:")
    print(param_counts)

    # --- Initialize optimizer ---
[content] train_dynamics.py (python)
args = tyro.cli(Args)


def dynamics_loss_fn(
    model: Genie, inputs: dict
) -> tuple[jax.Array, tuple[jax.Array | None, dict]]:
    """Compute masked dynamics loss"""
    model.train()
    outputs = model(inputs, training=True)
    mask = outputs["mask"]
    outputs["token_logits"] = outputs["token_logits"].astype(jnp.float32)
    ce_loss = optax.softmax_cross_entropy_with_integer_labels(
        outputs["token_logits"], outputs["video_tokens"]
    )
    ce_loss = (mask * ce_loss).sum() / mask.sum()
    acc = outputs["token_logits"].argmax(-1) == outputs["video_tokens"]
    acc = (mask * acc).sum() / mask.sum()
    select_probs = jax.nn.softmax(outputs["token_logits"])
    _, index_counts_tokenizer = jnp.unique_counts(
        jnp.ravel(outputs["video_tokens"]), size=args.num_patch_latents, fill_value=0
    )
    codebook_usage_tokenizer = (index_counts_tokenizer != 0).mean()
    metrics = dict(
        cross_entropy_loss=ce_loss,
        masked_token_accuracy=acc,
        select_logit=outputs["token_logits"].max(-1).mean(),
        select_p=select_probs.max(-1).mean(),
        entropy=jax.scipy.special.entr(select_probs).sum(-1).mean(),
        codebook_usage_tokenizer=codebook_usage_tokenizer,
    )
    return ce_loss, (None, metrics)


# @nnx.jit
def train_step(
    model: Genie, optimizer: nnx.Optimizer, inputs: dict
) -> tuple[jax.Array, jax.Array | None, dict]:
    """Update state and compute metrics"""

    def loss_fn(model: Genie) -> tuple[jax.Array, tuple[jax.Array | None, dict]]:
        return dynamics_loss_fn(model, inputs)

    (loss, (recon, metrics)), grads = nnx.value_and_grad(loss_fn, has_aux=True)(model)
    optimizer.update(grads)
    if args.log_gradients:
        metrics["gradients_std/"] = jax.tree.map(
            lambda x: x.std(), grads["params"]["dynamics"]
        )
    return loss, recon, metrics


if __name__ == "__main__":
    jax.distributed.initialize()
    num_devices = jax.device_count()
    if num_devices == 0:
        raise ValueError("No JAX devices found.")
    print(f"Running on {num_devices} devices.")

    if args.batch_size % num_devices != 0:
        raise ValueError(
            f"Global batch size {args.batch_size} must be divisible by "
            f"number of devices {num_devices}."
        )

    per_device_batch_size_for_init = args.batch_size // num_devices

    rng = jax.random.key(args.seed)

    # --- Initialize model ---
[content] train_tokenizer.py (python)
    jax.config.update("jax_transfer_guard", "disallow")
    while step < args.num_steps:
        for videos in dataloader:
            # --- Train step ---
            inputs = dict(videos=videos)
            loss, recon, metrics = train_step(tokenizer, optimizer, inputs)
            # metrics["lr"] = lr_schedule(step)
            with jax.default_device('cpu'):
                step_on_cpu = jnp.array(step)
                metrics["lr"] = float(lr_schedule(step_on_cpu))
[content] train_tokenizer.py (python)
.wandb_id:
            wandb_init_kwargs.update(
                {
                    "id": args.wandb_id,
                    "resume": "allow",
                }
            )
        wandb.init(**wandb_init_kwargs)

        wandb.config.update({"model_param_count": param_counts})

    print("Parameter counts:")
    print(param_counts)

    # --- Initialize optimizer ---
[content] train_tokenizer.py (python)
 on {num_devices} devices.")

    if args.batch_size % num_devices != 0:
        raise ValueError(
            f"Global batch size {args.batch_size} must be divisible by "
            f"number of devices {num_devices}."
[content] train_tokenizer.py (python)
tree.map(
            lambda x: x.std(), grads["params"]["encoder"]
        )
        metrics["vq_gradients_std/"] = jax.tree.map(
            lambda x: x.std(), grads["params"]["vq"]
        )
        metrics["decoder_gradients_std/"] = jax.tree.map(
            lambda x: x.std(), grads["params"]["decoder"]
[content] train_tokenizer.py (python)
(pix.ssim(gt, recon)).mean()
    _, index_counts = jnp.unique_counts(
        jnp.ravel(outputs["indices"]), size=args.num_latents, fill_value=0
    )
    codebook_usage = (index_counts != 0).mean()
    metrics = dict(
        loss=loss,
        mse=mse,
        q_loss=q_loss,
        commitment_loss=commitment_loss,
        psnr=psnr,
        ssim=ssim,
        codebook_usage=codebook_usage,
    )
    return loss, (outputs["recon"], metrics)


@nnx.jit
def train_step(
    tokenizer: TokenizerVQVAE, optimizer: nnx.Optimizer, inputs: dict
[content] train_tokenizer.py (python)
 we at least train the tokenizer natively in int8?
    gt = jnp.asarray(inputs["videos"], dtype=jnp.float32) / 255.0
    inputs["videos"] = gt.astype(args.dtype)
    model.train()
    outputs = model(inputs, training=True)
    outputs["recon"] = outputs["recon"].astype(jnp.float32)
    mse = jnp.square(gt - outputs["recon"]).mean()
    q_loss = jnp.square(jax.lax.stop_gradient(outputs["emb"]) - outputs["z"]).mean()
    commitment_loss = jnp.square(
        outputs["emb"] - jax.lax.stop_gradient(outputs["z"])
    ).mean()
    loss = mse + q_loss + args.vq_beta * commitment_loss

    # --- Compute validation metrics ---
    gt = gt.clip(0, 1).reshape(-1, *gt.shape[2:])
    recon = outputs["recon"].clip(0, 1).reshape(-1, *outputs["recon"].shape[2:])
    psnr =
[content] train_tokenizer.py (python)

def build_dataloader(args: Args) -> tuple[grain.DataLoader, grain.DataLoaderIterator]:
    image_shape = (args.image_height, args.image_width, args.image_channels)
    array_record_files = [
        os.path.join(args.data_dir, x)
        for x in os.listdir(args.data_dir)
        if x.endswith(".array_record")
    ]
    grain_dataloader = get_dataloader(
        array_record_files,
        args.seq_len,
        # NOTE: We deliberately pass t