p-doom/crowd-code-dataset-0.2 huggingface dataset:
[content] genie.py (python)
        self, batch: Dict[str, jax.Array], training: bool = True
    ) -> Dict[str, jax.Array]:
        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)