[content] genie.py (python)
    handler_registry.add('dataloader_state', grain.checkpoint.CheckpointRestore, grain.checkpoint.CheckpointHandler)
    

    checkpoint_options = ocp.CheckpointManagerOptions(
        step_format_fixed_length=6,
    )
    tokenizer_checkpoint_manager = ocp.CheckpointManager(
        directory=args.tokenizer_checkpoint,
        options=checkpoint_options,
        handler_registry=handler_registry,
    )
    dummy_tokenizer = TokenizerVQVAE(
        in_dim=args.image_channels,
        model_dim=args.tokenizer_dim,
        ffn_dim=args.tokenizer_ffn_dim,
        latent_dim=args.latent_patch_dim,
        num_latents=args.num_patch_latents,
        patch_size=args.patch_size,
        num_blocks=args.tokenizer_num_blocks,
        num_heads=args.tokenizer_num_heads,
        dropout=args.dropout,
        codebook_dropout=args.dropout,
        param_dtype=args.param_dtype,
        dtype=args.dtype,
        use_flash_attention=args.use_flash_attention,
    )
    tokenizer_init_params = dummy_tokenizer.init(_rng, inputs)
    dummy_tokenizer_train_state = TrainState.create(
        apply_fn=dummy_tokenizer.apply, params=tokenizer_init_params, tx=dummy_tx
    )
    abstract_sharded_tokenizer_state = _create_abstract_sharded_pytree(
        dummy_tokenizer_train_state, sharding
    )
    restored_tokenizer = tokenizer_checkpoint_manager.restore(
        step=tokenizer_checkpoint_manager.latest_step(),
        args=ocp.args.Composite(
            model_state=ocp.args.StandardRestore(abstract_sharded_tokenizer_state),
            dataloader_state=grain.checkpoint.CheckpointRestore(grain_iterator),
        ),
    )["model_state"]
    restored_tokenizer_params = restored_tokenizer.params["params"]
    train_state.params["params"]["tokenizer"].update(restored_tokenizer_params)
    tokenizer_checkpoint_manager.close()

    if args.lam_checkpoint:
        lam_checkpoint_manager = ocp.CheckpointManager(
            directory=args.lam_checkpoint,
            options=checkpoint_options,
            handler_registry=handler_registry,
        )
        dummy_lam = LatentActionModel(
            in_dim=args.image_channels,
            model_dim=args.lam_dim,
            ffn_dim=args.lam_ffn_dim,
            latent_dim=args.latent_patch_dim,
            num_latents=args.num_latent_actions,
            patch_size=args.lam_patch_size,
            num_blocks=args.lam_num_blocks,
            num_heads=args.lam_num_heads,
            dropout=args.dropout,
            codebook_dropout=args.dropout,
            param_dtype=args.param_dtype,
            dtype=args.dtype,
            use_flash_attention=args.use_flash_attention,
        )
        lam_init_params = dummy_lam.init(_rng, inputs)
        dummy_lam_train_state = TrainState.create(
            apply_fn=dummy_lam.apply, params=lam_init_params, tx=dummy_tx
        )
        abstract_sharded_lam_state = _create_abstract_sharded_pytree(
            dummy_lam_train_state, sharding
        )
        restored_lam = lam_checkpoint_manager.restore(
            step=lam_checkpoint_manager.latest_step(),
            args=ocp.args.Composite(
                model_state=ocp.args.StandardRestore(abstract_sharded_lam_state),
                dataloader_state=grain.checkpoint.CheckpointRestore(grain_iterator),