[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))

    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))