p-doom/crowd-code-dataset-0.2 huggingface dataset:
[content] utils/nn.py (python)
        if self.has_variable('cache', 'cached_key'):
            # Reset cached key and value to zeros
            cached_key = self.variable('cache', 'cached_key', jnp.zeros, 
                                     self.get_variable('cache', 'cached_key').shape,
                                     self.get_variable('cache', 'cached_key').dtype)
            cached_value = self.variable('cache', 'cached_value', jnp.zeros,
                                       self.get_variable('cache', 'cached_value').shape,
                                       self.get_variable('cache', 'cached_value').dtype)
            cache_index = self.variable('cache', 'cache_index', lambda: jnp.array(0, dtype=jnp.int32))
            
            # Set the cache variables to their initial state
            cached_key.value = jnp.zeros_like(cached_key.value)
            cached_value.value = jnp.zeros_like(cached_value.value)
            cache_index.value = jnp.array(0, dtype=jnp.int32)
[content] utils/nn.py (python)
            # Get the current cache variables
            cached_key = self.get_variable('cache', 'cached_key')
            cached_value = self.get_variable('cache', 'cached_value')
            cache_index = self.get_variable('cache', 'cache_index')
            
            # Reset the cache variables to their initial state
            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))