[content] train_dynamics.py (python)
    val_iterator: Optional[grain.DataLoaderIterator],
    restore_step: Optional[int] = None,
) -> tuple[
    int, nnx.Optimizer, grain.DataLoaderIterator, grain.DataLoaderIterator, jax.Array
]:
    step = 0
    if restore_step is None:
        restore_step = checkpoint_manager.latest_step()
    if args.restore_ckpt:
        abstract_optimizer = nnx.eval_shape(lambda: optimizer)
        abstract_optimizer_state = nnx.state(abstract_optimizer)
        if args.val_data_dir:
            restore_args = ocp.args.Composite(
                model_state=ocp.args.PyTreeRestore(abstract_optimizer_state),  # type: ignore
                train_dataloader_state=grain.checkpoint.CheckpointRestore(train_iterator),  # type: ignore
                val_dataloader_state=grain.checkpoint.CheckpointRestore(val_iterator),  # type: ignore
            )
        else:
            restore_args = ocp.args.Composite(
                model_state=ocp.args.PyTreeRestore(abstract_optimizer_state),  # type: ignore
                train_dataloader_state=grain.checkpoint.CheckpointRestore(train_iterator),  # type: ignore
            )
        restored = checkpoint_manager.restore(
            checkpoint_manager.latest_step(), args=restore_args
        )
        restored_optimizer_state = restored["model_state"]
        nnx.update(optimizer, restored_optimizer_state)
        train_iterator = restored["train_dataloader_state"]
        if args.val_data_dir:
            val_iterator = restored["val_dataloader_state"]
        step = checkpoint_manager.latest_step() or 0
        print(f"Restored dataloader and model state from step {step}")
    else:
        # Restore from pre-trained tokenizer (and LAM)
        rng, _rng = jax.random.split(rng)
        optimizer = restore_genie_components(optimizer, replicated_sharding, _rng, args)
        # NOTE: We have to remove the (unused) tokenizer vq dropout due flax.nnx lazily initializing modules.
        # Specifically, the first dynamics model checkpoint will contain the vq dropout module,
        # but the first full restore will fail due to nnx not initializing the module when
        # dropout is set to 0.0.
        del optimizer.model.tokenizer.vq.drop
    return step, optimizer, train_iterator, val_iterator, rng


def main(args: Args) -> None:
    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}."
        )

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

    # --- Initialize model ---
    genie, rng = build_model(args, rng)
    _, params, _ = nnx.split(genie, nnx.Param, ...)
    param_counts = count_parameters_by_component(params)

    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 ---
    optimizer, lr_schedule = build_optimizer(genie, args)
    del genie

    # FIXME: switch to create_hybrid_device_mesh for runs spanning multiple nodes
    _, replicated_sharding, videos_sharding = build_mesh_and_sharding(num_devices)

    shard_optimizer_states(optimizer, replicated_sharding)

    # --- Initialize checkpoint manager ---
    checkpoint_manager = build_checkpoint_manager(args)

    # --- Create DataLoaderIterator from dataloader ---
    train_iterator = build_dataloader(args, args.data_dir)
    val_iterator = None
    if args.val_data_dir:
        val_iterator = build_dataloader(args, args.val_data_dir)

    # --- Restore checkpoint ---
    step, optimizer, train_iterator, val_iterator, rng = (
        restore_or_initialize_components(
            args,
            checkpoint_manager,
            optimizer,
            train_iterator,
            rng,
            replicated_sharding,
            val_iterator,
        )
    )

    # --- Define loss and train step (close over args) ---
    def dynamics_loss_fn(
        model: Genie,
        inputs: dict,
        training: bool = False,
        pred_full_frame: bool = False,
    ) -> tuple[jax.Array, tuple[jax.Array, dict]]:
        gt = jnp.asarray(inputs["videos"], dtype=jnp.float32) / 255.0
        inputs["videos"] = gt.astype(args.dtype)
        outputs = model(inputs, training=training, pred_full_frame=pred_full_frame)
        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"])
        gt_val = 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_val, recon)).mean()
        ssim = jnp.asarray(pix.ssim(gt_val, recon)).mean()
        _, index_counts_lam = jnp.unique_counts(
            jnp.ravel(outputs["lam_indices"]),
            size=args.num_latent_actions,
            fill_value=0,
        )
        _, index_counts_tokenizer = jnp.unique_counts(
            jnp.ravel(outputs["video_tokens"]),
            size=args.num_patch_latents,
            fill_value=0,
        )
        codebook_usage_lam = (index_counts_lam != 0).mean()
        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(),
            psnr=psnr,
            ssim=ssim,
            codebook_usage_lam=codebook_usage_lam,
            codebook_usage_tokenizer=codebook_usage_tokenizer,
        )
        return ce_loss, (outputs["recon"], metrics)

    @nnx.jit(donate_argnums=0)
    def train_step(
        optimizer: nnx.Optimizer, inputs: dict
    ) -> tuple[jax.Array, jax.Array, dict]:
        def loss_fn(model: Genie) -> tuple[jax.Array, tuple[jax.Array, dict]]:
            model.train()
            return dynamics_loss_fn(model, inputs, training=True)

        (loss, (recon, metrics)), grads = nnx.value_and_grad(loss_fn, has_aux=True)(
            optimizer.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

    @nnx.jit
    def val_step(
        genie: Genie, inputs: dict
    ) -> tuple[jax.Array, jax.Array, dict, jax.Array, jax.Array, dict]:
        """Evaluate model and compute metrics"""
        genie.eval()
        (loss, (recon, metrics)) = dynamics_loss_fn(
            genie, inputs, training=False, pred_full_frame=False
        )
        (loss_full_frame, (recon_full_frame, metrics_full_frame)) = dynamics_loss_fn(
            genie, inputs, training=False, pred_full_frame=True
        )
        return (
            loss,
            recon,
            metrics,
            loss_full_frame,
            recon_full_frame,
            metrics_full_frame,
        )

    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 = []
        inputs = None
        recon = None
        recon_full_frame = None
        for videos in val_dataloader:
            rng, _rng_mask = jax.random.split(rng, 2)
            inputs = dict(videos=videos, mask_rng=_rng_mask)
            (
                loss,
                recon,
                metrics,
                loss_full_frame,
                recon_full_frame,
                metrics_full_frame,
            ) = val_step(genie, inputs)
            loss_per_step.append(loss)
            metrics_per_step.append(metrics)
            loss_full_frame_per_step.append(loss_full_frame)
            metrics_full_frame_per_step.append(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_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_losses = {
            "val_loss": np.mean(loss_per_step),
            "val_loss_full_frame": np.mean(loss_full_frame_per_step),
        }
        val_metrics.update(val_metrics_full_frame)
        val_metrics.update(val_losses)
        return val_metrics, inputs, recon, recon_full_frame

    # --- TRAIN LOOP ---
    dataloader_train = (
        jax.make_array_from_process_local_data(videos_sharding, elem)
        for elem in train_iterator
    )
    dataloader_val = None
    if val_iterator:
        dataloader_val = (
            jax.make_array_from_process_local_data(videos_sharding, elem)
            for elem in val_iterator
        )
    if jax.process_index() == 0:
        first_videos = next(dataloader_train)
        sample_inputs = dict(videos=first_videos, mask_rng=rng)
        compiled = train_step.lower(optimizer, sample_inputs).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_videos], dataloader_train)
    print(f"Starting training from step {step}...")
    first_step = step
    while step < args.num_steps:
        for videos in dataloader_train:
            # --- Train step ---
            rng, _rng_mask = jax.random.split(rng, 2)
            inputs = dict(videos=videos, mask_rng=_rng_mask)
            loss, recon, metrics = train_step(optimizer, inputs)
            if step == first_step:
                print_mem_stats("After params initialized")
            metrics["lr"] = lr_schedule(step)
            print(f"Step {step}, loss: {loss}")
            step += 1

            # --- Validation loss ---
            val_results = {}
            if dataloader_val and step % args.val_interval == 0:
                rng, _rng_mask_val = jax.random.split(rng, 2)
                print("Calculating validation metrics...")
                val_metrics, val_gt_batch, val_recon, val_full_frame = (
                    calculate_validation_metrics(
                        dataloader_val, optimizer.model, _rng_mask_val
                    )
                )
                print(f"Step {step}, validation loss: {val_metrics['val_loss']}")
                val_results = {
                    "metrics": val_metrics,
                    "gt_batch": val_gt_batch,
                    "recon": val_recon,
                    "full_frame": val_full_frame,
                }

            # --- Logging ---
            if args.log:
                if step % args.log_interval == 0 and jax.process_index() == 0:
                    log_dict = {"loss": loss, "step": step, **metrics}
                    if val_results:
                        log_dict.update(val_results["metrics"])
                    wandb.log(log_dict)
                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 val_results:
                        val_results["gt_seq_val"] = (
                            val_results["gt_batch"]["videos"][0].astype(jnp.float32)
                            / 255.0
                        )
                        val_results["recon_seq_val"] = val_results["recon"].clip(0, 1)
                        val_comparison_seq = jnp.concatenate(
                            (val_results["gt_seq_val"], val_results["recon_seq_val"]),
                            axis=1,
                        )
                        val_results["val_comparison_seq"] = einops.rearrange(
                            val_comparison_seq * 255, "t h w c -> h (t w) c"
                        )
                        val_results["full_frame_seq_val"] = val_results["full_frame"][
                            0
                        ].clip(0, 1)
                        val_results["val_full_frame_comparison_seq"] = jnp.concatenate(
                            (
                                val_results["gt_seq_val"],
                                val_results["full_frame_seq_val"],
                            ),
                            axis=1,
                        )
                        val_results["val_full_frame_comparison_seq"] = einops.rearrange(
                            val_results["val_full_frame_comparison_seq"] * 255, "t h w c -> h (t w) c"
                        )
                    # NOTE: Process-dependent control flow deliberately happens
                    # after indexing operation since it must not contain code
                    # sections that lead to cross-accelerator communication.
                    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))
                            ),
                        )
                        if val_results:
                            log_images.update(
                                dict(
                                    val_image=wandb.Image(
                                        np.asarray(
                                            val_results["gt_seq_val"][args.seq_len - 1]
                                        )
                                    ),
                                    val_recon=wandb.Image(
                                        np.asarray(
                                            val_results["recon_seq_val"][
                                                args.seq_len - 1
                                            ]
                                        )
                                    ),
                                    val_true_vs_recon=wandb.Image(
                                        np.asarray(
                                            val_results["val_comparison_seq"].astype(
                                                np.uint8
                                            )
                                        )
                                    ),
                                    val_full_frame=wandb.Image(
                                        np.asarray(
                                            val_results["full_frame_seq_val"][
                                                args.seq_len - 1
                                            ]
                                        )
                                    ),
                                    val_true_vs_full_frame=wandb.Image(
                                        np.asarray(
                                            val_results[
                                                "val_full_frame_comparison_seq"
                                            ].astype(np.uint8)
                                        )
                                    ),
                                )
                            )
                        wandb.log(log_images)
            # --- Checkpointing ---
            if args.save_ckpt and step % args.log_checkpoint_interval == 0:
                optimizer_state = nnx.state(optimizer)
                if args.val_data_dir:
                    ckpt_manager_args = ocp.args.Composite(
                        model_state=ocp.args.PyTreeSave(optimizer_state),  # type: ignore
                        train_dataloader_state=grain.checkpoint.CheckpointSave(  # type: ignore
                            train_iterator  # type: ignore
                        ),
                        val_dataloader_state=grain.checkpoint.CheckpointSave(  # type: ignore
                            val_iterator  # type: ignore
                        ),
                    )
                else:
                    ckpt_manager_args = ocp.args.Composite(
                        model_state=ocp.args.PyTreeSave(optimizer_state),  # type: ignore
                        train_dataloader_state=grain.checkpoint.CheckpointSave(  # type: ignore
                            train_iterator  # type: ignore
                        ),
                    )
                checkpoint_manager.save(step, args=ckpt_manager_args)