[content] jasmine/train_dynamics.py (python)
    if args.dyna_type == "diffusion":
        if "x_pred" in outputs.keys():
            mse_x = jnp.mean((outputs["x_pred"] - outputs["x_gt"]) ** 2, axis=(1, 2, 3))
            mse = jnp.mean(mse_x)
        else:
            mse = 0.0  # TODO mihir: fix this
    else:
        mask_BTN = 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_BTN * ce_loss).sum() / mask_BTN.sum()
        z_val = jax.nn.logsumexp(outputs["token_logits"], axis=-1)
        z_loss_metric = (mask_BTN * (z_val**2)).sum() / mask_BTN.sum()

        masked_token_top_1_acc = _calculate_top_k_accuracy(
            outputs["token_logits"], outputs["video_tokens"], mask_BTN, 1
        )
        masked_token_top_2_acc = _calculate_top_k_accuracy(
            outputs["token_logits"], outputs["video_tokens"], mask_BTN, 2
        )
        masked_token_top_5_acc = _calculate_top_k_accuracy(
            outputs["token_logits"], outputs["video_tokens"], mask_BTN, 5
        )
        masked_token_top_16_acc = _calculate_top_k_accuracy(
            outputs["token_logits"], outputs["video_tokens"], mask_BTN, 16
        )

        select_probs = jax.nn.softmax(outputs["token_logits"])
        _, index_counts_tokenizer = jnp.unique_counts(
            jnp.ravel(outputs["video_tokens"]),
            size=num_patch_latents,
            fill_value=0,
        )
        codebook_usage_tokenizer = (index_counts_tokenizer != 0).mean()
    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()
    metrics = dict(
        psnr=psnr,
        ssim=ssim,
    )
    if args.dyna_type == "diffusion":
        metrics["mse"] = mse
        loss = mse
    else:
        metrics["cross_entropy_loss"] = ce_loss
        metrics["masked_token_top1_accuracy"] = masked_token_top_1_acc
        metrics["masked_token_top2_accuracy"] = masked_token_top_2_acc
        metrics["masked_token_top5_accuracy"] = masked_token_top_5_acc
        metrics["masked_token_top16_accuracy"] = masked_token_top_16_acc
        metrics["select_logit"] = outputs["token_logits"].max(-1).mean()
        metrics["select_p"] = select_probs.max(-1).mean()
        metrics["entropy"] = jax.scipy.special.entr(select_probs).sum(-1).mean()
        metrics["z_loss"] = z_loss_metric
        metrics["codebook_usage_tokenizer"] = codebook_usage_tokenizer
        loss = ce_loss
    if "lam_indices" in outputs.keys():
        _, index_counts_lam = jnp.unique_counts(
            jnp.ravel(outputs["lam_indices"]),
            size=num_actions,
            fill_value=0,
        )
        codebook_usage_lam = (index_counts_lam != 0).mean()
        metrics["codebook_usage_lam"] = codebook_usage_lam
    return loss, metrics


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

    # FIXME: switch to create_hybrid_device_mesh for runs spanning multiple nodes
    _, replicated_sharding, videos_sharding, actions_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,
    ) -> 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)
        loss, metrics = _calculate_step_metrics(
            outputs, gt, args.num_actions, args.num_patch_latents
        )
        if "z_loss" in metrics.keys():
            z_loss = metrics["z_loss"]
            total_loss = loss + args.z_loss_weight * z_loss
        else: