[content] train_dynamics.py (python)
            lam_indices = genie.vq_encode(inputs, training=False)
            tokenizer_outputs = genie.tokenizer.vq_encode(
                inputs["videos"], training=False
            )
            tokens_full_frame = tokenizer_outputs["indices"]
            inputs["latent_actions"] = lam_indices
            gt = jnp.asarray(inputs["videos"], dtype=jnp.float32) / 255.0
            inputs["videos"] = gt[:, :-1].astype(
                args.dtype
            )  # remove last frame for generation
            recon_full_frame, logits_full_frame = genie.sample(
                inputs,
                args.seq_len,
                args.val_temperature,
                args.val_sample_argmax,
                args.val_maskgit_steps,
            )
            step_outputs = {
                "recon": recon_full_frame,
                "token_logits": logits_full_frame,
                "video_tokens": tokens_full_frame,
                "mask": jnp.zeros_like(tokens_full_frame).at[:, -1].set(True),
                "lam_indices": lam_indices,
            }
            loss_full_frame, metrics_full_frame = _calculate_step_metrics(
                step_outputs, gt, args.num_actions, args.num_patch_latents
            )
            val_output.update(
                {
                    "loss_full_frame": loss_full_frame,
                    "recon_full_frame": recon_full_frame,
                    "metrics_full_frame": metrics_full_frame,
                }
            )
        return val_output

    def calculate_validation_metrics(val_dataloader, genie, rng):
        step = 0
        loss_per_step = []
        metrics_per_step = []
        loss_full_frame_per_step = []
        metrics_full_frame_per_step = []
        batch = None
        recon = None
        recon_full_frame = None
        for batch in val_dataloader:
            rng, _rng_mask = jax.random.split(rng, 2)
            batch["rng"] = _rng_mask
            val_outputs = val_step(genie, batch)
            loss_per_step.append(val_outputs["loss"])
            metrics_per_step.append(val_outputs["metrics"])
            recon = val_outputs["recon"]
            if args.eval_full_frame:
                loss_full_frame_per_step.append(val_outputs["loss_full_frame"])
                metrics_full_frame_per_step.append(val_outputs["metrics_full_frame"])
                recon_full_frame = val_outputs["recon_full_frame"]
            step += 1
            if step > args.val_steps:
                break

        if step < args.val_steps:
            print(
                f"Warning: Your validation dataset is too small to make val_steps many steps. Made {step} steps, expected {args.val_steps}"
            )

        val_metrics = {
            f"val_{key}": np.mean([float(m[key]) for m in metrics_per_step])
            for key in metrics_per_step[0].keys()
        }
        val_metrics["val_loss"] = np.mean(loss_per_step)
        if args.eval_full_frame:
            val_metrics_full_frame = {
                f"val_full_frame_{key}": np.mean(
                    [float(m[key]) for m in metrics_full_frame_per_step]
                )
                for key in metrics_full_frame_per_step[0].keys()
            }
            val_metrics.update(val_metrics_full_frame)
            val_metrics["val_loss_full_frame"] = np.mean(loss_full_frame_per_step)
        return val_metrics, batch, recon, recon_full_frame

    # --- TRAIN LOOP ---
    dataloader_train = (
        {
            "videos": jax.make_array_from_process_local_data(
                videos_sharding, local_data=elem["videos"]
            ),
            "actions": (
                jax.make_array_from_process_local_data(
                    actions_sharding, elem["actions"]
                )
                if args.use_gt_actions
                else None
            ),
        }
        for elem in train_iterator
    )
    dataloader_val = None
    if val_iterator:
        dataloader_val = (
            {
                "videos": jax.make_array_from_process_local_data(
                    videos_sharding, elem["videos"]
                ),
                "actions": (
                    jax.make_array_from_process_local_data(
                        actions_sharding, elem["actions"]
                    )
                    if args.use_gt_actions
                    else None
                ),
            }
            for elem in val_iterator
        )
    if jax.process_index() == 0:
        first_batch = next(dataloader_train)
        first_batch["rng"] = rng  # type: ignore
        compiled = train_step.lower(optimizer, first_batch).compile()
        print_compiled_memory_stats(compiled.memory_analysis())
        print_compiled_cost_analysis(compiled.cost_analysis())
        # Do not skip the first batch during training
        dataloader_train = itertools.chain([first_batch], dataloader_train)
    print(f"Starting training from step {step}...")
    first_step = step
    while step < args.num_steps:
        for batch in dataloader_train:
            # --- Train step ---
            rng, _rng_mask = jax.random.split(rng, 2)
            batch["rng"] = _rng_mask
            loss, recon, metrics = train_step(optimizer, batch)
            if step == first_step:
                print_mem_stats("After params initialized")
            # metrics["lr"] = lr_schedule(step)
            # print(f"Step {step}, loss: {loss}")