[content] train_dynamics.py (python)
    @nnx.jit
    def val_step(genie: Genie, inputs: dict) -> dict:
        """Evaluate model and compute metrics"""
        genie.eval()
        (loss, (recon, metrics)) = dynamics_loss_fn(genie, inputs, training=False)
        val_output = {"loss": loss, "recon": recon, "metrics": metrics}

        # --- Evaluate full frame prediction (sampling) ---
        if args.eval_full_frame:
            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_latent_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"])
            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"])
            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_dataloader, val_metrics, batch, recon, recon_full_frame