[content] sample.py (python)
def _reset_spatial_kv_cache(variables):
    """
    Reset spatial KV cache for all STBlocks in the dynamics transformer.
    This function manipulates the cache variables directly without changing module structure.
    """
    # Get the dynamics cache
    dynamics_cache = variables["cache"]["dynamics"]
    
    # Reset cache for each STBlock
    for i in range(args.dyna_num_blocks):
        block_name = f"st_block_{i}"
        if block_name in dynamics_cache:
            spatial_attn_cache = dynamics_cache[block_name]["spatial_attn"]["cache"]
            
            # Reset cached_key, cached_value, and cache_index
            if "cached_key" in spatial_attn_cache:
                cached_key = spatial_attn_cache["cached_key"]
                spatial_attn_cache["cached_key"] = jnp.zeros_like(cached_key)
            
            if "cached_value" in spatial_attn_cache:
                cached_value = spatial_attn_cache["cached_value"]
                spatial_attn_cache["cached_value"] = jnp.zeros_like(cached_value)
            
            if "cache_index" in spatial_attn_cache:
                spatial_attn_cache["cache_index"] = jnp.array(0, dtype=jnp.int32)
    
    return variables

def _custom_sample_causal_with_kv_reset(batch, seq_len, temperature, sample_argmax):
    """
    Custom implementation of sample_causal that resets spatial KV cache between frames.
    This avoids modifying the module structure while achieving the same effect.
    """
    # Encode context frames
    tokenizer_out = inference_genie.apply(
        variables_for_sampling["params"], 
        dict(videos=batch["videos"]), 
        False, 
        method=inference_genie.tokenizer.vq_encode
    )
    token_idxs = tokenizer_out["indices"]  # (B, T, N)
    B, T, N = token_idxs.shape
    
    # Prepare initial token sequence
    pad_shape = (B, seq_len - T, N)
    token_idxs_full = jnp.concatenate(
        [token_idxs, jnp.zeros(pad_shape, dtype=token_idxs.dtype)], axis=1
    )  # (B, seq_len, N)
    
    # Prepare latent actions
    action_tokens = inference_genie.apply(
        variables_for_sampling["params"],
        dict(latent_actions=batch["latent_actions"]),
        False,
        method=lambda params, batch: inference_genie.lam.vq.get_codes(batch["latent_actions"])
    )
    
    # Autoregressive generation loop
    rng = batch["rng"]
    current_variables = variables_for_sampling.copy()
    
    for t in range(T, seq_len):
        for n in range(N):
            # Run dynamics to get logits
            dyna_inputs = {
                "video_tokens": token_idxs_full,
                "latent_actions": action_tokens
            }
            dyna_outputs = inference_genie.apply(
                current_variables["params"],
                dyna_inputs,
                False,
                method=inference_genie.dynamics
            )
            
            # Get logits for current position
            next_token_logits = dyna_outputs["token_logits"][:, t, n, :].astype(jnp.float32)
            
            # 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
                )
            
            # Update token sequence
            token_idxs_full = token_idxs_full.at[:, t, n].set(next_token)
        
        # Reset spatial KV cache after completing all patches for a frame
        # We need to reset the cache by reinitializing the variables
        current_variables = unfreeze(current_variables)
        # Reset the dynamics cache to initial state
        if "dynamics" in current_variables["cache"]:
            # Reinitialize the dynamics cache
            rng, _rng = jax.random.split(rng)
            dummy_dyna_inputs = {
                "video_tokens": jnp.zeros((B, 1, N), dtype=token_idxs.dtype),
                "latent_actions": jnp.zeros((B, 1, 1), dtype=action_tokens.dtype)
            }
            _, new_cache = inference_genie.dynamics.init_with_output(_rng, dummy_dyna_inputs, training=False)
            current_variables["cache"]["dynamics"] = new_cache["dynamics"]
        current_variables = freeze(current_variables)
    
    # Decode final tokens
    final_frames = inference_genie.apply(
        current_variables["params"],
        dict(indices=token_idxs_full, video_hw=batch["videos"].shape[2:4]),
        False,
        method=inference_genie.tokenizer.decode
    )
    
    return final_frames