p-doom/crowd-code-dataset-0.2 huggingface dataset:
[content] genie.py (python)
        videos_BTHWC = batch["videos"]
        tokenizer_outputs = self.tokenizer.vq_encode(videos_BTHWC, training=False)
        token_indices_BTN = tokenizer_outputs["indices"]
        lam_outputs = self.lam.vq_encode(videos_BTHWC, training=False)
        z_q_BTm11L = lam_outputs["z_q"]
        action_indices_E = lam_outputs["indices"]
        latent_actions_BTm11L = jax.lax.cond(
            self.lam_co_train,
            lambda: z_q_BTm11L,
            lambda: jax.lax.stop_gradient(z_q_BTm11L),
        )
        outputs = dict(
            video_tokens=jax.lax.stop_gradient(token_indices_BTN),
            latent_actions=latent_actions_BTm11L,
        )
        outputs["mask_rng"] = batch["mask_rng"]
        dyna_logits_BTNV, dyna_mask = self.dynamics(outputs, training)
        outputs["token_logits"] = dyna_logits_BTNV
        if dyna_mask is not None:
            outputs["mask"] = dyna_mask
        mle_indices_BTN = jnp.argmax(outputs["token_logits"], axis=-1)
        H, W = batch["videos"].shape[2:4]
        outputs["recon"] = self.tokenizer.decode(mle_indices_BTN, (H, W))
        outputs["lam_indices"] = action_indices_E