[content] /fast/home/franz.srambical/jafar/genie.py (python)
        # FIXME (f.srambical): reset spatial kv cache after each frame
        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)
        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], step_n: jax.Array
        ) -> tuple[tuple[jax.Array, jax.Array, jax.Array, jax.Array], None]:
            rng, token_idxs_BSN, 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)
            step_t += 1

            new_carry = (rng, token_idxs_BSN, action_tokens_EL, step_t)
            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)

            # --- Reset spatial KV caches before each frame ---
            for block in dynamics_causal.transformer.blocks:
                assert block.spatial_attention.cached_key is not None
                assert block.spatial_attention.cached_value is not None
                assert block.spatial_attention.cache_index.value is not None
                block.spatial_attention.cached_key.value = jnp.zeros((B * T, 1, self.dyna_dim), dtype=self.dtype)
                block.spatial_attention.cached_value.value = jnp.zeros((B * T, 1, self.dyna_dim), dtype=self.dtype)
                block.temporal_attention.cache_index.value = jnp.array(0, dtype=jnp.int32)