[content] train_tokenizer.py (python)

    return step, optimizer, grain_iterator


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 ---
    tokenizer, rng = build_model(args, rng)

    _, params, _ = nnx.split(tokenizer, 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(tokenizer, args)

    # FIXME: switch to create_hybrid_device_mesh for runs spanning multiple nodes
    mesh, 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 ---
    grain_dataloader, grain_iterator = build_dataloader(args)

    # --- Restore checkpoint ---
    step, optimizer, grain_iterator = restore_checkpoint_if_needed(
        args, checkpoint_manager, optimizer, grain_iterator
    )

    # --- 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