[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
        return outputs

    def sample(
        self,
        batch: Dict[str, jax.Array],
        seq_len: int,
        steps: int = 25,
        temperature: float = 1,
        sample_argmax: bool = False,
    ) -> jax.Array:
        """
        Autoregressively samples up to `seq_len` future frames, following Figure 8 of the paper.

        - Input frames are tokenized once.
        - Future frames are generated autoregressively in token space.
        - All frames are detokenized in a single pass.

        Note:
        - For interactive or step-wise sampling, detokenization should occur after each action.
        - To maintain consistent tensor shapes across timesteps, all current and future frames are decoded at every step.
        - Temporal causal structure is preserved by
            a) reapplying the mask before each decoding step.
            b) a temporal causal mask is applied within each ST-transformer block.

        Dimension keys:
            B: batch size
            T: number of input (conditioning) frames
            N: number of patches per frame
            M: model dimension
            S: sequence length
            H: height
            W: width
            D: B * T * N
            E: B * (T - 1)
        """
        # --- Encode videos and actions ---
        videos_BTHWC = batch["videos"]
        latent_actions_E = batch["latent_actions"]
        tokenizer_out = self.tokenizer.vq_encode(videos_BTHWC, training=False)
        token_idxs_BTN = tokenizer_out["indices"]
        B, T, N = token_idxs_BTN.shape
        pad_shape = (B, seq_len - T, N)
        pad = jnp.zeros(pad_shape, dtype=token_idxs_BTN.dtype)
        token_idxs_BSN = jnp.concatenate([token_idxs_BTN, pad], axis=1)
        action_tokens_EL = self.lam.vq.get_codes(latent_actions_E)

        def maskgit_step_fn(
            carry: tuple[jax.Array, jax.Array, jax.Array, jax.Array], step: jax.Array
        ) -> tuple[tuple[jax.Array, jax.Array, jax.Array, jax.Array], None]:
            rng, token_idxs_BSN, mask_BSN, action_tokens_EL = carry
            S, N = token_idxs_BSN.shape[1:]

            # --- Construct + encode video ---
            vid_embed_BSNM = self.dynamics.patch_embed(token_idxs_BSN)
            mask_token_111M = self.dynamics.mask_token.value
            mask_expanded_BSN1 = mask_BSN[..., None]
            vid_embed_BSNM = jnp.where(mask_expanded_BSN1, mask_token_111M, vid_embed_BSNM)

            # --- Predict transition ---
            action_tokens_BSm1L = jnp.reshape(action_tokens_EL, (B, S - 1, 1))
            act_embed_BSm1M = self.dynamics.action_up(action_tokens_BSm1L)
            # FIXME (f.srambical): We must not pad the actions, but remove the last frame (https://github.com/p-doom/jasmine/issues/122)
            vid_embed_BSNM += jnp.pad(act_embed_BSm1M, ((0, 0), (1, 0), (0, 0), (0, 0)))
            unmasked_ratio = jnp.cos(jnp.pi * (step + 1) / (steps * 2))
            step_temp = temperature * (1.0 - unmasked_ratio)
            final_logits_BSNV = self.dynamics.transformer(vid_embed_BSNM) / step_temp

            # --- Sample new tokens for final frame ---
            if sample_argmax:
                sampled_token_idxs_BSN = jnp.argmax(final_logits_BSNV, axis=-1)
            else:
                rng, _rng = jax.random.split(rng)
                sampled_token_idxs_BSN = jax.random.categorical(_rng, final_logits_BSNV)
            gather_fn = jax.vmap(jax.vmap(jax.vmap(lambda x, y: x[y])))
            final_token_probs_BSN = gather_fn(
                jax.nn.softmax(final_logits_BSNV), sampled_token_idxs_BSN
            )
            final_token_probs_BSN += ~mask_BSN
            # Update masked tokens only
            token_idxs_BSN = jnp.where(mask_BSN, sampled_token_idxs_BSN, token_idxs_BSN)

            # --- Update mask ---
            num_unmasked_tokens = jnp.round(N * (1.0 - unmasked_ratio)).astype(int)
            idx_mask_N = jnp.arange(final_token_probs_BSN.shape[-1]) > num_unmasked_tokens
            sorted_idxs_BSN = jnp.argsort(final_token_probs_BSN, axis=-1, descending=True)
            mask_update_fn = jax.vmap(lambda msk, ids: msk.at[ids].set(idx_mask_N))
            new_mask_BSN = mask_update_fn(mask_BSN, sorted_idxs_BSN)

            new_carry = (rng, token_idxs_BSN, new_mask_BSN, action_tokens_EL)
            return new_carry, None

        def generation_step_fn(
            carry: tuple[jax.Array, jax.Array], step_t: jax.Array
        ) -> tuple[tuple[jax.Array, jax.Array], None]:
            rng, current_token_idxs_BSN = carry
            rng, step_rng = jax.random.split(rng)

            # Mask current and future frames (i.e., t >= step_t)
            mask_S = jnp.arange(seq_len) >= step_t
            mask_BSN = jnp.broadcast_to(mask_S[None, :, None], (B, seq_len, N)).astype(
                bool
            )
            masked_token_idxs_BSN = current_token_idxs_BSN * ~mask_BSN

            # --- Initialize and run MaskGIT loop ---
            init_carry_maskgit = (
                step_rng,
                masked_token_idxs_BSN,
                mask_BSN,
                action_tokens_EL,
            )
            final_carry_maskgit, _ = jax.lax.scan(
                maskgit_step_fn, init_carry_maskgit, jnp.arange(steps)
            )
            updated_token_idxs = final_carry_maskgit[1]
            new_carry = (rng, updated_token_idxs)
            return new_carry, None

        # --- Run the autoregressive generation using jax.lax.scan ---
        initial_carry = (batch["rng"], token_idxs_BSN)
        timesteps_to_scan = jnp.arange(T, seq_len)
        final_carry, _ = jax.lax.scan(
            generation_step_fn, initial_carry, timesteps_to_scan
        )
        final_token_idxs = final_carry[1]

        # --- Decode all tokens at once at the end ---
        H, W = batch["videos"].shape[2:4]
        final_frames = self.tokenizer.decode(
            final_token_idxs,
            video_hw=(H, W),
        )
        return final_frames

    def vq_encode(self, batch: Dict[str, jax.Array], training: bool) -> jax.Array:
        # --- Preprocess videos ---
        video_BTHWC = batch["videos"]
        lam_output = self.lam.vq_encode(video_BTHWC, training=training)
        lam_indices_E = lam_output["indices"]
        return lam_indices_E