[content] /hkfs/home/project/hk-project-p0023960/tum_cte0515/Projects/jasmine/genie.py (python)
            masked_logits_BSNV = current_logits_BSNV * jnp.expand_dims(~mask_BSN, -1)

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

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

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

    def sample_causal(
        self,
        batch: Dict[str, jax.Array],
        seq_len: int,
        temperature: float = 1,
        sample_argmax: bool = False,
    ) -> tuple[jax.Array, jax.Array, 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
            E: B * (S - 1)
        """
        assert isinstance(self.dynamics, DynamicsCausal)
        # --- 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)
        logits_BSNV = jnp.zeros((*token_idxs_BSN.shape, self.num_patch_latents))
        action_tokens_EL = self.lam.vq.get_codes(latent_actions_E)
        dynamics_causal: DynamicsCausal = self.dynamics

        def causal_step_fn(
            carry: tuple[jax.Array, jax.Array, jax.Array, jax.Array, jax.Array], step_n: jax.Array
        ) -> tuple[tuple[jax.Array, jax.Array, jax.Array, jax.Array, jax.Array], None]:
            rng, token_idxs_BSN, logits_BSNV, action_tokens_EL, step_t = carry
            S, N = token_idxs_BSN.shape[1:]
            L = action_tokens_EL.shape[-1]

            # --- Construct + encode video ---
            vid_embed_BSNM = dynamics_causal.patch_embed(token_idxs_BSN)

            # --- Predict transition ---
            action_tokens_BSm1L = jnp.reshape(action_tokens_EL, (B, S - 1, L))
            act_embed_BSm1M = dynamics_causal.action_up(action_tokens_BSm1L)
            act_embed_BSM = jnp.pad(act_embed_BSm1M, ((0, 0), (1, 0), (0, 0)))
            act_embed_BS1M = jnp.reshape(
                act_embed_BSM, (B, S, 1, act_embed_BSM.shape[-1])
            )
            vid_embed_BSNp1M = jnp.concatenate([act_embed_BS1M, vid_embed_BSNM], axis=2)
            final_logits_BTNp1V = (
                dynamics_causal.transformer(vid_embed_BSNp1M, (step_t, step_n))
                / temperature
            )
            final_logits_BV = final_logits_BTNp1V[:, step_t, step_n, :]

            # --- Sample new tokens for final frame ---
            if sample_argmax:
                sampled_token_idxs_B = jnp.argmax(final_logits_BV, axis=-1)
            else:
                rng, _rng = jax.random.split(rng)
                sampled_token_idxs_B = jax.random.categorical(_rng, final_logits_BV)
            # Update next tokens only
            token_idxs_BSN = token_idxs_BSN.at[:, step_t, step_n].set(
                sampled_token_idxs_B
            )
            logits_BSNV = logits_BSNV.at[:, step_t, step_n].set(final_logits_BV)