[content] genie.py (python)
    def sample_causal(
        self,
        batch: Dict[str, Any],
        seq_len: int,
        temperature: float = 1.0,
        sample_argmax: bool = False,
    ):
        """
        Autoregressively samples with KV caching.
        """
        # This method assumes the model was initialized with decode=True
        assert self.decode, "sample_causal with KV cache requires model to be in decode mode."

        # --- Encode context frames ---
        tokenizer_out = self.tokenizer.vq_encode(batch["videos"], training=False)
        token_idxs = tokenizer_out["indices"]  # (B, T, N)
        B, T, N = token_idxs.shape

        # --- Get latent actions ---
        action_tokens = self.lam.vq.get_codes(batch["latent_actions"]) # (B, S-1, D)

        # --- Initialize cache with context frames ---
        # The first call to the dynamics model will populate the KV cache.
        # We only care about the logits for the very last token of the context
        # to predict the first new token.
        dyna_inputs = {
            "video_tokens": token_idxs,
            "latent_actions": action_tokens[:, :T-1] 
        }
        # Note: We are now calling self.dynamics directly
        logits = self.dynamics(dyna_inputs, training=False)["token_logits"]
        
        # Logits for the first token of the first generated frame (t=T)
        next_token_logits = logits[:, T-1, N-1, :]
        
        # The full sequence of tokens to be filled in
        generated_tokens = jnp.zeros((B, seq_len, N), dtype=token_idxs.dtype)
        generated_tokens = generated_tokens.at[:, :T].set(token_idxs)

        # --- Autoregressive generation loop ---
        rng = batch["rng"]
        
        # We have N tokens per frame. We start generating from frame T.
        # The total number of tokens to generate is (seq_len - T) * N
        for t_step in range(T * N, seq_len * N):
            frame_idx = t_step // N
            patch_idx = t_step % N

            # Sample or argmax
            if sample_argmax:
                next_token = jnp.argmax(next_token_logits, axis=-1)
            else:
                rng, step_rng = jax.random.split(rng)
                next_token = jax.random.categorical(
                    step_rng, next_token_logits / temperature, axis=-1
                )
            
            # Place the new token in the sequence
            generated_tokens = generated_tokens.at[:, frame_idx, patch_idx].set(next_token)

            # If we just generated the last token, break
            if frame_idx == seq_len - 1 and patch_idx == N - 1:
                break
            
            # Prepare input for the next step: just the token we generated
            # The model will use its cache for all previous tokens.
            next_input_token = next_token[:, None, None] # Reshape to (B, 1, 1)
            
            # Get the action for the current frame
            # Action for frame `i` influences transition from `i-1` to `i`
            current_action = action_tokens[:, frame_idx-1:frame_idx] # (B, 1, D)

            next_dyna_inputs = {
                "video_tokens": next_input_token,
                "latent_actions": current_action,
            }

            # This call uses and updates the cache, and is very fast!
            next_token_logits = self.dynamics(next_dyna_inputs, training=False)["token_logits"]
            next_token_logits = next_token_logits[:, 0, 0, :] # Squeeze singleton dims

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