p-doom/crowd-code-dataset-0.2 huggingface dataset:
[content] utils/nn.py (python)
        # Reset spatial attention cache
        if self.has_variable('spatial_attention', 'cached_key'):
            cached_key = self.get_variable('spatial_attention', 'cached_key')
            cached_value = self.get_variable('spatial_attention', 'cached_value')
            self.put_variable('spatial_attention', 'cached_key', jnp.zeros_like(cached_key))
            self.put_variable('spatial_attention', 'cached_value', jnp.zeros_like(cached_value))
            self.put_variable('spatial_attention', 'cache_index', jnp.array(0, dtype=jnp.int32))
[content] utils/nn.py (python)
class MultiHeadAttentionWithKVCacheReset(nn.MultiHeadAttention):
    """
    MultiHeadAttention module with a reset_kv_cache() method to reset the key-value cache.
    """
    def reset_kv_cache(self):
        """
        Resets the key-value cache, if present.
        """
        if self.has_variable('cache', 'cached_key'):
            cached_key = self.get_variable('cache', 'cached_key')
            cached_value = self.get_variable('cache', 'cached_value')
            
            self.put_variable('cache', 'cached_key', jnp.zeros_like(cached_key))
            self.put_variable('cache', 'cached_value', jnp.zeros_like(cached_value))
            self.put_variable('cache', 'cache_index', jnp.array(0, dtype=jnp.int32))
[content] utils/nn.py (python)
        # Apply reset to all blocks
        jax.lax.scan(
            lambda carry, block_idx: (carry, reset_block_cache(block_idx)),
            init=None,
            xs=jnp.arange(self.num_blocks)
        )
[content] utils/nn.py (python)
        """
        Reset KV cache for all transformer blocks using scan.
        """
        def reset_block_cache(block_idx):
            block = STBlock(
                dim=self.model_dim,
                num_heads=self.num_heads,
                dropout=self.dropout,
                param_dtype=self.param_dtype,
                dtype=self.dtype,
                spatial_bert=self.spatial_bert,
                use_flash_attention=self.use_flash_attention,
                use_kv_cache=self.use_kv_cache,
                name=f"block_{block_idx}"
            )
            block.reset_kv_cache()
[content] utils/nn.py (python)
                name=f"block_{block_idx}"
            )
        
        # Apply blocks using scan
        def apply_block(carry, block_idx):
            block = create_block(block_idx)
            return block(carry), None
        
        x, _ = nn.scan(
            apply_block,
            variable_axes={"params": 0},  # Each block has its own parameters
            split_rngs={"params": True},  # Each block gets its own RNG
            length=self.num_blocks,
            unroll=1  # Unroll for better performance with small num_blocks
        )(x, jnp.arange(self.num_blocks))
[content] utils/nn.py (python)
            block.reset_kv_cache()
        
        # Apply reset to all blocks
        jax.lax.scan(
            lambda carry, block_idx: (carry, reset_block_cache(block_idx)),
            init=None,
            xs=jnp.arange(self.num_blocks)
        )
[content] utils/nn.py (python)
        """
        Reset KV cache for all transformer blocks.
        """
        for i in range(self.num_blocks):
            block = self.get_variable('params', f'block_{i}')
            # Access the block module and reset its cache
            block_module = STBlock(
                dim=self.model_dim,
                num_heads=self.num_heads,
                dropout=self.dropout,
                param_dtype=self.param_dtype,
                dtype=self.dtype,
                spatial_bert=self.spatial_bert,
                use_flash_attention=self.use_flash_attention,
                use_kv_cache=self.use_kv_cache,
                name=f"block_{i}"
            )
            block_module.reset_kv_cache()
[content] utils/nn.py (python)
        """
        Resets the key-value cache for all STBlock instances in the transformer.
        """
        # Get all STBlock instances and reset their KV caches
        for i in range(self.num_blocks):
            st_block = self.get_variable('STBlock_' + str(i), 'STBlock')
            if hasattr(st_block, 'reset_kv_cache'):
                st_block.reset_kv_cache()
[content] utils/nn.py (python)
    def setup(self):
        # Create STBlock instances that can be accessed later for cache reset
        self.st_blocks = []
        for i in range(self.num_blocks):
            self.st_blocks.append(
                STBlock(
                    dim=self.model_dim,
                    num_heads=self.num_heads,
                    dropout=self.dropout,
                    param_dtype=self.param_dtype,
                    dtype=self.dtype,
                    spatial_bert=self.spatial_bert,
                    use_flash_attention=self.use_flash_attention,
                    use_kv_cache=self.use_kv_cache,
                )
            )
[content] utils/nn.py (python)
        for _ in range(self.num_blocks):
            x = STBlock(
                dim=self.model_dim,
                num_heads=self.num_heads,
                dropout=self.dropout,
                param_dtype=self.param_dtype,
                dtype=self.dtype,
                spatial_bert=self.spatial_bert,
                use_flash_attention=self.use_flash_attention,
                use_kv_cache=self.use_kv_cache,
            )(x)
[content] utils/nn.py (python)
self.spatial_attn = MultiHeadAttentionWithKVCacheReset(
            num_heads=self.num_heads,
            qkv_features=self.dim,
            dropout_rate=self.dropout,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )
[content] utils/nn.py (python)
MultiHeadAttentionWithKVCacheReset(
            num_heads=self.num_heads,
            qkv_features=self.dim,
            dropout_rate=self.dropout,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
            attention_fn=_create_flash_attention_fn(self.use_flash_attention, is_causal=not self.spatial_bert),
            decode=self.use_kv_cache,
        )
[content] utils/nn.py (python)
nn.MultiHeadAttention(
            num_heads=self.num_heads,
            qkv_features=self.dim,
            dropout_rate=self.dropout,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
            attention_fn=_create_flash_attention_fn(self.use_flash_attention, is_causal=True),
            decode=self.use_kv_cache,
            name="temporal_attn",
        # FIXME (f.srambical): check whether we should still pass the mask if we set is_causal=True
        )
[content] utils/nn.py (python)
        
        # Create output projection
        self.output_projection = nn.Dense(
            self.out_dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )

    def __call__(self, x: jax.Array) -> jax.Array:
        # Input projection and normalization
        x = self.input_norm1(x)
        x = self.input_projection(x)
        x = self.input_norm2(x)
        
        # Process through ST blocks
[content] utils/nn.py (python)
        # Create input projection layers
        self.input_norm1 = nn.LayerNorm(
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )
        self.input_projection = nn.Dense(
            self.model_dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )
        self.input_norm2 = nn.LayerNorm(
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )
        
        # Create ST blocks
[content] utils/nn.py (python)
        
        # Create ST blocks with explicit names
        for i in range(self.num_blocks):
            setattr(self, f'st_block_{i}', STBlock(
                dim=self.model_dim,
                num_heads=self.num_heads,
                dropout=self.dropout,
                param_dtype=self.param_dtype,
                dtype=self.dtype,
                spatial_bert=self.spatial_bert,
                use_flash_attention=self.use_flash_attention,
                use_kv_cache=self.use_kv_cache,
            ))
[content] utils/nn.py (python)
        self.st_blocks: List[STBlock] = []
        for _ in range(self.num_blocks):
            self.st_blocks.append(
                STBlock(
                    dim=self.model_dim,
                    num_heads=self.num_heads,
                    dropout=self.dropout,
                    param_dtype=self.param_dtype,
                    dtype=self.dtype,
                    spatial_bert=self.spatial_bert,
                    use_flash_attention=self.use_flash_attention,
                    use_kv_cache=self.use_kv_cache,
                )
            )
[content] utils/nn.py (python)
# Create input projection layers
        self.input_norm1 = nn.LayerNorm(
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )
        self.input_projection = nn.Dense(
            self.model_dim,
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )
        self.input_norm2 = nn.LayerNorm(
            param_dtype=self.param_dtype,
            dtype=self.dtype,
        )
        
        # Create ST blocks
[content] utils/nn.py (python)
    @nn.compact
    def __call__(self, x: jax.Array) -> jax.Array:
        x = nn.Sequential(
            [
                nn.LayerNorm(
                    param_dtype=self.param_dtype,
                    dtype=self.dtype,
                ),
                nn.Dense(self.model_dim,
                param_dtype=self.param_dtype,
                dtype=self.dtype,
                ),
                nn.LayerNorm(
                    param_dtype=self.param_dtype,
                    dtype=self.dtype,
                ),
            ]
        )(x)
[content] utils/nn.py (python)
            x = STBlock(
                dim=self.model_dim,
                num_heads=self.num_heads,
                dropout=self.dropout,
                param_dtype=self.param_dtype,
                dtype=self.dtype,
                spatial_bert=self.spatial_bert,
                use_flash_attention=self.use_flash_attention,
                use_kv_cache=self.use_kv_cache,
            )(x)
        x = nn.Dense(
[content] utils/nn.py (python)
= self.st_blocks[i](x)
        
        x = self.output_projection(x)
        return x
    
    def reset_spatial_kv_cache(self):
        """
        Resets the spatial key-value cache for all STBlocks in the transformer.
        """
        for st_block in self.st_blocks:
            st_block.reset_spatial_kv_cache(
[content] utils/nn.py (python)
        class MultiHeadAttentionWithKVCacheReset(nn.MultiHeadAttention):
    """
    MultiHeadAttention module with a reset_kv_cache() method to reset the key-value cache.
    """
    def reset_kv_cache(self):
        """
        Resets the key-value cache, if present.
        """
        if self.has_variable('cache', 'cached_key'):
            cached_key = self.get_variable('cache', 'cached_key')
            cached_value = self.get_variable('cache', 'cached_value')
            
            self.put_variable('cache', 'cached_key', jnp.zeros_like(cached_key))
            self.put_variable('cache', 'cached_value', jnp.zeros_like(cached_value))
            self.put_variable('cache', 'cache_index', jnp.array(0, dtype=jnp.int32))
[content] utils/nn.py (python)
    def reset_spatial_kv_cache(self):
        """
        Resets the spatial key-value cache for all STBlocks in the transformer.
        """
        for st_block in self.st_blocks:
            st_block.reset_spatial_kv_cache()
[content] sample.py (python)
# --- Define autoregressive sampling loop with external KV cache resetting ---
def _autoreg_sample_with_kv_reset(rng, video_batch, action_batch):
    vid = video_batch[:, : args.start_frame + 1]
    sampling_fn = nn.apply(_sampling_wrapper, inference_genie)
    
    # Initialize variables for sampling
    current_variables = variables_for_sampling.copy()
    
    # Get tokenizer output for initial frames
    tokenizer_out = inference_genie.apply(
        current_variables["params"], 
        dict(videos=vid), 
        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, args.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(
        current_variables["params"],
        dict(latent_actions=action_batch),
        False,
        method=lambda params, batch: inference_genie.lam.vq.get_codes(batch["latent_actions"])
    )
    
    # Autoregressive generation loop
    rng, _rng = jax.random.split(rng)
    for t in range(T, args.seq_len):
        for n in range(N):
            # Create batch for this step
            batch = dict(
                videos=vid, 
                latent_actions=action_batch, 
                rng=_rng,
                # Pass the current token sequence
                _token_idxs_full=token_idxs_full,
                _current_t=t,
                _current_n=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 args.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 / args.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
        current_variables = _reset_spatial_kv_cache(current_variables)
    
    # Decode final tokens
    final_frames = inference_genie.apply(
        current_variables["params"],
        dict(indices=token_idxs_full, video_hw=vid.shape[2:4]),
        False,
        method=inference_genie.tokenizer.decode
    )
    
    return final_frames
[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
[content] genie.py (python)
    def reset_spatial_kv_cache_most_direct(self):
        """
        The most direct way to reset spatial KV cache - accessing cache variables
        directly from the MultiHeadAttention modules without any wrapper methods.
        This shows the exact cache variable names and structure.
        """
        # Access the dynamics STTransformer
        dynamics_transformer = self.dynamics.dynamics
        
        # Iterate through all STBlocks
        for i in range(dynamics_transformer.num_blocks):
            # Get the STBlock by name
            st_block = getattr(dynamics_transformer, f'st_block_{i}')
            
            # Access the spatial attention module directly
            spatial_attn = st_block.spatial_attn
            
            # Check if cache variables exist
            if spatial_attn.has_variable('cache', 'cached_key'):
                # Get current cache values
                cached_key = spatial_attn.get_variable('cache', 'cached_key')
                cached_value = spatial_attn.get_variable('cache', 'cached_value')
                cache_index = spatial_attn.get_variable('cache', 'cache_index')
                
                # Reset to zero/initial state
                spatial_attn.put_variable('cache', 'cached_key', jnp.zeros_like(cached_key))
                spatial_attn.put_variable('cache', 'cached_value', jnp.zeros_like(cached_value))
                spatial_attn.put_variable('cache', 'cache_index', jnp.array(0, dtype=jnp.int32))
[content] genie.py (python)
    def reset_spatial_kv_cache_direct(self):
        """
        Directly resets the spatial key-value cache for all STBlocks in the dynamics model
        by accessing the cache variables without using the wrapper methods.
        """
        for i in range(self.dynamics.dynamics.num_blocks):
            st_block = getattr(self.dynamics.dynamics, f'st_block_{i}')
            spatial_attn = st_block.spatial_attn
            if spatial_attn.has_variable('cache', 'cached_key'):
                cached_key = spatial_attn.get_variable('cache', 'cached_key')
                cached_value = spatial_attn.get_variable('cache', 'cached_value')
                
                spatial_attn.put_variable('cache', 'cached_key', jnp.zeros_like(cached_key))
                spatial_attn.put_variable('cache', 'cached_value', jnp.zeros_like(cached_value))
                spatial_attn.put_variable('cache', 'cache_index', jnp.array(0, dtype=jnp.int32))

    def reset_spatial_kv_cache_most_direct(self):
        """
        The most direct way to reset spatial KV cache - accessing cache variables
        directly from the MultiHeadAttention modules without any wrapper methods.
        This shows the exact cache variable names and structure.
        """
        # Access the dynamics STTransformer
        dynamics_transformer = self.dynamics.dynamics
        
        # Iterate through all STBlocks
        for i in range(dynamics_transformer.num_blocks):
            # Get the STBlock by name
            st_block = getattr(dynamics_transformer, f'st_block_{i}')
            
            # Access the spatial attention module directly
            spatial_attn = st_block.spatial_attn
            
            # Check if cache variables exist
            if spatial_attn.has_variable('cache', 'cached_key'):
                # Get current cache values
                cached_key = spatial_attn.get_variable('cache', 'cached_key')
                cached_value = spatial_attn.get_variable('cache', 'cached_value')
                cache_index = spatial_attn.get_variable('cache', 'cache_index')
                
                # Reset to zero/initial state
                spatial_attn.put_variable('cache', 'cached_key', jnp.zeros_like(cached_key))
                spatial_attn.put_variable('cache', 'cached_value', jnp.zeros_like(cached_value))
                spatial_attn.put_variable('cache', 'cache_index', jnp.array(0, dtype=jnp.int32))
[content] genie.py (python)
    def reset_all_kv_cache_direct(self):
        """
        Demonstrates the most direct way to reset KV cache for both spatial and temporal attention.
        This shows how to access cache variables directly from both wrapper and base MultiHeadAttention classes.
        """
        dynamics_transformer = self.dynamics.dynamics
        
        for i in range(dynamics_transformer.num_blocks):
            st_block = getattr(dynamics_transformer, f'st_block_{i}')
            
            # Reset spatial attention cache (uses wrapper)
            spatial_attn = st_block.spatial_attn
            if spatial_attn.has_variable('cache', 'cached_key'):
                cached_key = spatial_attn.get_variable('cache', 'cached_key')
                cached_value = spatial_attn.get_variable('cache', 'cached_value')
                
                spatial_attn.put_variable('cache', 'cached_key', jnp.zeros_like(cached_key))
                spatial_attn.put_variable('cache', 'cached_value', jnp.zeros_like(cached_value))
                spatial_attn.put_variable('cache', 'cache_index', jnp.array(0, dtype=jnp.int32))
            
            # Reset temporal attention cache (uses base nn.MultiHeadAttention)
            temporal_attn = st_block.temporal_attn
            if temporal_attn.has_variable('cache', 'cached_key'):
                cached_key = temporal_attn.get_variable('cache', 'cached_key')
                cached_value = temporal_attn.get_variable('cache', 'cached_value')
                
                temporal_attn.put_variable('cache', 'cached_key', jnp.zeros_like(cached_key))
                temporal_attn.put_variable('cache', 'cached_value', jnp.zeros_like(cached_value))
                temporal_attn.put_variable('cache', 'cache_index', jnp.array(0, dtype=jnp.int32))
[content] genie.py (python)
        
        # Reset spatial KV cache directly by accessing cache variables
        # This is the most direct way without using any wrapper methods
        for i in range(self.dynamics.dynamics.num_blocks):
            st_block = getattr(self.dynamics.dynamics, f'st_block_{i}')
            spatial_attn = st_block.spatial_attn
            if spatial_attn.has_variable('cache', 'cached_key'):
                cached_key = spatial_attn.get_variable('cache', 'cached_key')
                cached_value = spatial_attn.get_variable('cache', 'cached_value')
                
                spatial_attn.put_variable('cache', 'cached_key', jnp.zeros_like(cached_key))
                spatial_attn.put_variable('cache', 'cached_value', jnp.zeros_like(cached_value))
                spatial_attn.put_variable('cache', 'cache_index', jnp.array(0, dtype=jnp.int32))
[content] genie.py (python)
            st_block = getattr(self.dynamics.dynamics, f'STBlock_{i}')
            # In STBlock, the first MultiHeadAttention is spatial (index 0)
            spatial_attn = getattr(st_block, 'MultiHeadAttention_0')
[content] genie.py (python)
        
        # Reset spatial KV cache directly by accessing cache variables
        # This is the most direct way without using any wrapper methods
        for i in range(self.dynamics.dynamics.num_blocks):
            st_block = getattr(self.dynamics.dynamics, f'STBlock_{i}')
            # In STBlock, the first MultiHeadAttention is spatial (index 0)
            spatial_attn = getattr(st_block, 'MultiHeadAttention_0')
            if spatial_attn.has_variable('cache', 'cached_key'):
                cached_key = spatial_attn.get_variable('cache', 'cached_key')
                cached_value = spatial_attn.get_variable('cache', 'cached_value')
                
                spatial_attn.put_variable('cache', 'cached_key', jnp.zeros_like(cached_key))
                spatial_attn.put_variable('cache', 'cached_value', jnp.zeros_like(cached_value))
                spatial_attn.put_variable('cache', 'cache_index', jnp.array(0, dtype=jnp.int32))
[content] genie.py (python)
    def reset_spatial_kv_cache_direct(self):
        """
        Directly resets the spatial key-value cache for all STBlocks in the dynamics model
        by accessing the cache variables without using the wrapper methods.
        """
        for i in range(self.dynamics.dynamics.num_blocks):
            st_block = getattr(self.dynamics.dynamics, f'STBlock_{i}')
            # In STBlock, the first MultiHeadAttention is spatial (index 0)
            spatial_attn = getattr(st_block, 'MultiHeadAttention_0')
            if spatial_attn.has_variable('cache', 'cached_key'):
                cached_key = spatial_attn.get_variable('cache', 'cached_key')
                cached_value = spatial_attn.get_variable('cache', 'cached_value')
                
                spatial_attn.put_variable('cache', 'cached_key', jnp.zeros_like(cached_key))
                spatial_attn.put_variable('cache', 'cached_value', jnp.zeros_like(cached_value))
                spatial_attn.put_variable('cache', 'cache_index', jnp.array(0, dtype=jnp.int32))

    def reset_spatial_kv_cache_most_direct(self):
        """
        The most direct way to reset spatial KV cache - accessing cache variables
        directly from the MultiHeadAttention modules without any wrapper methods.
        This shows the exact cache variable names and structure.
        """
        # Access the dynamics STTransformer
        dynamics_transformer = self.dynamics.dynamics
        
        # Iterate through all STBlocks
        for i in range(dynamics_transformer.num_blocks):
            # Get the STBlock by default name
            st_block = getattr(dynamics_transformer, f'STBlock_{i}')
            
            # Access the spatial attention module directly (first MultiHeadAttention)
            spatial_attn = getattr(st_block, 'MultiHeadAttention_0')
            
            # Check if cache variables exist
            if spatial_attn.has_variable('cache', 'cached_key'):
                # Get current cache values
                cached_key = spatial_attn.get_variable('cache', 'cached_key')
                cached_value = spatial_attn.get_variable('cache', 'cached_value')
                cache_index = spatial_attn.get_variable('cache', 'cache_index')
                
                # Reset to zero/initial state
                spatial_attn.put_variable('cache', 'cached_key', jnp.zeros_like(cached_key))
                spatial_attn.put_variable('cache', 'cached_value', jnp.zeros_like(cached_value))
                spatial_attn.put_variable('cache', 'cache_index', jnp.array(0, dtype=jnp.int32))
[content] genie.py (python)
    def reset_all_kv_cache_direct(self):
        """
        Demonstrates the most direct way to reset KV cache for both spatial and temporal attention.
        This shows how to access cache variables directly using Flax's default module names.
        """
        dynamics_transformer = self.dynamics.dynamics
        
        for i in range(dynamics_transformer.num_blocks):
            st_block = getattr(dynamics_transformer, f'STBlock_{i}')
            
            # Reset spatial attention cache (first MultiHeadAttention)
            spatial_attn = getattr(st_block, 'MultiHeadAttention_0')
            if spatial_attn.has_variable('cache', 'cached_key'):
                cached_key = spatial_attn.get_variable('cache', 'cached_key')
                cached_value = spatial_attn.get_variable('cache', 'cached_value')
                
                spatial_attn.put_variable('cache', 'cached_key', jnp.zeros_like(cached_key))
                spatial_attn.put_variable('cache', 'cached_value', jnp.zeros_like(cached_value))
                spatial_attn.put_variable('cache', 'cache_index', jnp.array(0, dtype=jnp.int32))
            
            # Reset temporal attention cache (second MultiHeadAttention)
            temporal_attn = getattr(st_block, 'MultiHeadAttention_1')
            if temporal_attn.has_variable('cache', 'cached_key'):
                cached_key = temporal_attn.get_variable('cache', 'cached_key')
                cached_value = temporal_attn.get_variable('cache', 'cached_value')
                
                temporal_attn.put_variable('cache', 'cached_key', jnp.zeros_like(cached_key))
                temporal_attn.put_variable('cache', 'cached_value', jnp.zeros_like(cached_value))
                temporal_attn.put_variable('cache', 'cache_index', jnp.array(0, dtype=jnp.int32))
[content] genie.py (python)
        # Get all variables from the dynamics transformer to find the cache variables
        variables = dynamics_transformer.variables
        
        # Iterate through all cache variables and reset them
        for collection_name, collection_vars in variables.items():
            if collection_name == 'cache':
                for var_name, var_value in collection_vars.items():
                    if 'cached_key' in var_name or 'cached_value' in var_name:
                        # Reset to zero
                        dynamics_transformer.put_variable(collection_name, var_name, jnp.zeros_like(var_value))
                    elif 'cache_index' in var_name:
                        # Reset index to 0
                        dynamics_transformer.put_variable(collection_name, var_name, jnp.array(0, dtype=jnp.int32))
[content] genie.py (python)
        # Iterate through all STBlocks
        for i in range(dynamics_transformer.num_blocks):
            # Get the STBlock by default name
            st_block = getattr(dynamics_transformer, f'STBlock_{i}')
            
            # Access the spatial attention module directly (first MultiHeadAttention)
            spatial_attn = getattr(st_block, 'MultiHeadAttention_0')
            
            # Check if cache variables exist
            if spatial_attn.has_variable('cache', 'cached_key'):
                # Get current cache values
                cached_key = spatial_attn.get_variable('cache', 'cached_key')
                cached_value = spatial_attn.get_variable('cache', 'cached_value')
                
                # Reset to zero/initial state
                spatial_attn.put_variable('cache', 'cached_key', jnp.zeros_like(cached_key))
                spatial_attn.put_variable('cache', 'cached_value', jnp.zeros_like(cached_value))
                spatial_attn.put_variable('cache', 'cache_index', jnp.array(0, dtype=jnp.int32))
[content] genie.py (python)
        The most direct way to reset spatial KV cache - accessing cache variables
        directly from the MultiHeadAttention modules without any wrapper methods.
        This shows the exact cache variable names and structure.