p-doom/crowd-code-dataset-0.2 huggingface dataset:
[content] sample.py (python)
def reset_spatial_cache_in_vars(variables: dict) -> dict:
    """
    Creates a new variables dictionary with the spatial attention cache reset.
    """
    if 'cache' not in variables:
        return variables
    cache_unfrozen = unfreeze(variables['cache'])
    if 'dynamics' in cache_unfrozen and 'dynamics' in cache_unfrozen['dynamics']:
        st_transformer_cache = cache_unfrozen['dynamics']['dynamics']
        for block_name in st_transformer_cache:
            if 'STBlock' in block_name and 'MultiHeadAttention_0' in st_transformer_cache[block_name]:
                spatial_attn_cache = st_transformer_cache[block_name]['MultiHeadAttention_0']
                if 'cache_index' in spatial_attn_cache:
                    spatial_attn_cache['cache_index'] = jnp.array(0, dtype=jnp.int32)
                if 'cached_key' in spatial_attn_cache:
                    spatial_attn_cache['cached_key'] = jnp.zeros_like(spatial_attn_cache['cached_key'])
                if 'cached_value' in spatial_attn_cache:
                    spatial_attn_cache['cached_value'] = jnp.zeros_like(spatial_attn_cache['cached_value'])
    return freeze({'params': variables['params'], 'cache': cache_unfrozen})
[content] sample.py (python)
def reset_spatial_cache_in_vars(variables: Mapping[str, Any]) -> Mapping[str, Any]:
    """
    Creates a new variables dictionary with the spatial attention cache reset.
    """
    if 'cache' not in variables:
        return variables
    cache_unfrozen = unfreeze(variables['cache'])
    if 'dynamics' in cache_unfrozen and 'dynamics' in cache_unfrozen['dynamics']:
        st_transformer_cache = cache_unfrozen['dynamics']['dynamics']
        for block_name in st_transformer_cache:
            if 'STBlock' in block_name and 'MultiHeadAttention_0' in st_transformer_cache[block_name]:
                spatial_attn_cache = st_transformer_cache[block_name]['MultiHeadAttention_0']
                if 'cache_index' in spatial_attn_cache:
                    spatial_attn_cache['cache_index'] = jnp.array(0, dtype=jnp.int32)
                if 'cached_key' in spatial_attn_cache:
                    spatial_attn_cache['cached_key'] = jnp.zeros_like(spatial_attn_cache['cached_key'])
                if 'cached_value' in spatial_attn_cache:
                    spatial_attn_cache['cached_value'] = jnp.zeros_like(spatial_attn_cache['cached_value'])
    return freeze({'params': variables['params'], 'cache': cache_unfrozen})
[content] sample.py (python)


def _sampling_wrapper(module, batch):
    # return module.sample_maskgit(batch, args.seq_len, args.maskgit_steps, args.temperature, args.sample_argmax)
    return module.sample_causal(batch, args.seq_len, args.temperature, args.sample_argmax)
[tab] models/tokenizer.py (python)
from typing import Dict, Any, Tuple

import flax.linen as nn
import jax.numpy as jnp

from utils.preprocess import patchify, unpatchify
from utils.nn import STTransformer, VectorQuantizer


class TokenizerVQVAE(nn.Module):
    """ST-ViVit VQ-VAE"""

    in_dim: int
    model_dim: int
    latent_dim: int
    num_latents: int
    patch_size: int
    num_blocks: int
    num_heads: int
    dropout: float
    codebook_dropout: float
    param_dtype: jnp.dtype
    dtype: jnp.dtype
    use_flash_attention: bool

    def setup(self):
        self.encoder = STTransformer(
            self.model_dim,
            self.latent_dim,
            self.num_blocks,
            self.num_heads,
            self.dropout,
            self.param_dtype,
            self.dtype,
            use_flash_attention=self.use_flash_attention,
            use_kv_cache=False,
        )
        self.vq = VectorQuantizer(
            self.latent_dim,
            self.num_latents,
            self.codebook_dropout,
        )
        self.out_dim = self.in_dim * self.patch_size**2
        self.decoder = STTransformer(
            self.model_dim,
            self.out_dim,
            self.num_blocks,
            self.num_heads,
            self.dropout,
            self.param_dtype,
            self.dtype,
            use_flash_attention=self.use_flash_attention,
            use_kv_cache=False,
        )

    def __call__(self, batch: Dict[str, Any], training: bool = True) -> Dict[str, Any]:
        H, W = batch["videos"].shape[2:4]
        outputs = self.vq_encode(batch["videos"], training)
        recon = self.decoder(outputs["z_q"])  # (B, T, H_down * W_down, C)
        recon = recon.astype(jnp.float32)
        recon = nn.sigmoid(recon)
        recon = recon.astype(self.dtype)
        outputs["recon"] = unpatchify(recon, self.patch_size, H, W)
        return outputs

    def vq_encode(self, videos: Any, training: bool = True) -> Dict[str, Any]:
        # --- Preprocess + encode ---
        B, T = videos.shape[:2]
        x = patchify(videos, self.patch_size)
        N = x.shape[2]
        x = self.encoder(x)  # (B, T, N, E)

        # --- Vector quantize ---
        x = x.reshape(B * T * N, self.latent_dim)
        z_q, z, emb, indices = self.vq(x, training)
        z_q = z_q.reshape(B, T, N, self.latent_dim)
        indices = indices.reshape(B, T, N)
        return dict(z_q=z_q, z=z, emb=emb, indices=indices)

    def decode(self, indices: Any, video_hw: Tuple[int, int]):
        z = self.vq.codebook[indices]
        recon = self.decoder(z)
        recon = recon.astype(jnp.float32)
        recon = nn.sigmoid(recon)
        recon = recon.astype(self.dtype)
        return unpatchify(recon, self.patch_size, *video_hw)
[content] sample.py (python)
action_batch_output = inference_genie.apply(variables_for_sampling["params"], batch, False, method=Genie.vq_encode)
# Extract indices from the output tuple
action_batch = action_batch_output.reshape(video_batch.shape[0], args.seq_len - 1, 1)
[content] sample.py (python)
action_batch_output = inference_genie.apply(variables_for_sampling["params"], batch, False, method=Genie.vq_encode)
# Extract indices from the output dictionary
action_batch = action_batch_output.reshape(video_batch.shape[0], args.seq_len - 1, 1)
[content] sample.py (python)
action_batch_output, _ = inference_genie.apply(variables_for_sampling["params"], batch, False, method=Genie.vq_encode)
# Extract indices from the output (it's already the indices array)
action_batch = action_batch_output.reshape(video_batch.shape[0], args.seq_len - 1, 1)
[content] genie.py (python)


    def sample_causal(
        self,
        batch: Dict[str, Any],
        seq_len: int,
        temperature: float = 1.0,
        sample_argmax: bool = False,
    ):
        """
        Autoregressively samples up to `seq_len` future frames using the causal transformer backend.

        - Input frames are tokenized once.
        - Future frames are generated one at a time, each conditioned on all previous frames.
        - All frames are detokenized in a single pass at the end.

        Args:
            batch: Dict with at least "videos" (B, T, H, W, C)
            seq_len: total number of frames to generate (including context)
            temperature: sampling temperature
            sample_argmax: if True, use argmax instead of sampling

        Returns:
            Generated video frames (B, seq_len, H, W, C)
        """
        # --- Encode context frames ---
        tokenizer_out = self.tokenizer.vq_encode(batch["videos"], training=False)
        token_idxs = tokenizer_out["indices"]  # (B, T, N)
        B, T, N = token_idxs.shape

        # jax.debug.print("token_idxs shape: {}", token_idxs.shape)
        # --- Prepare initial token sequence ---
        # Pad with zeros for future frames
        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 = self.lam.vq.get_codes(batch["latent_actions"]) # (B, S-1, )
        # --- Autoregressive generation loop ---
        rng = batch["rng"]
        for t in range(T, seq_len):
            for n in range(N):
                jax.debug.print("Sampling token {} from frame {}", n, t)
                dyna_inputs = {
                    "video_tokens": token_idxs_full,
                    "latent_actions": action_tokens
                    }
                # jax.debug.print("token_idxs_full 0: {}", token_idxs_full[0,:,0])
                dyna_outputs = self.dynamics(dyna_inputs, training=False)
                # # dyna_outputs["token_logits"]: (B, t, N, vocab_size)
                # # We want the logits for the last time step (frame t-1 predicting t)
                # jax.debug.breakpoint()
                next_token_logits = dyna_outputs["token_logits"][:, t, n, :].astype(jnp.float32)  # (B, 1, vocab_size)

                # Sample or argmax for each patch
                if sample_argmax:
                    next_token = jnp.argmax(next_token_logits, axis=-1)  # (B, 1)
                else:
                    rng, step_rng = jax.random.split(rng)
                    next_token = jax.random.categorical(
                        step_rng, next_token_logits / temperature, axis=-1
                    )  # (B, 1)

                # Insert the generated tokens into the sequence
                token_idxs_full = token_idxs_full.at[:, t, n].set(next_token)

        # --- Decode all tokens at once at the end ---
        final_frames = self.tokenizer.decode(
            token_idxs_full, video_hw=batch["videos"].shape[2:4]
        )
        return final_frames
[content] sample.py (python)
    # --- 1. Define the JIT-compiled step function ---
    # @jax.jit
    def predict_token_fn(variables: Mapping[str, Any], inputs: Mapping[str, Any]) -> tuple[Mapping[str, Any], Mapping[str, Any]]:
        """
        This function calls the dynamics submodule's apply method.
        """
        return inference_genie.apply(
            variables,
            inputs,
            method = lambda module, inputs : module.dynamics(inputs, training=False),
            mutable=['cache']
        )

    # --- 2. Initialization ---
    variables = variables_for_sampling
    context_frames = video_batch[:, : args.start_frame + 1]

    tokenizer_params = {'tokenizer': variables['params']['tokenizer']}
    tokenizer_out = inference_genie.apply(
        {'params': tokenizer_params},
        context_frames,
        method=lambda module, videos: module.tokenizer.vq_encode(videos, training=False)
    )
    context_tokens = tokenizer_out['indices']
    B, T_context, N = context_tokens.shape
    
    token_sequence = jnp.zeros((B, args.seq_len, N), dtype=context_tokens.dtype)
    token_sequence = token_sequence.at[:, :T_context, :].set(context_tokens)

    lam_params = {'lam': variables['params']['lam']}
    action_codes = inference_genie.apply(
        {'params': lam_params},
        action_batch,
        method=lambda module, actions: module.lam.vq.get_codes(actions)
    )

    # --- 3. The Sampling Loop (State Management) ---
    for t in range(T_context, args.seq_len):
        print(f"Generating frame {t}... resetting spatial cache.")
        variables = reset_spatial_cache_in_vars(variables)

        for n in range(N):
            step_inputs = {
                "video_tokens": token_sequence,
                "latent_actions": action_codes,
            }
            
            dynamics_variables = {'params': variables['params']['dynamics'], 'cache': variables['cache']['dynamics']}
            
            dyna_outputs, updated_dynamics_cache = predict_token_fn(dynamics_variables, step_inputs)
            
            variables_unfrozen = unfreeze(variables)
            variables_unfrozen['cache']['dynamics'] = updated_dynamics_cache['cache']['dynamics']
            variables = freeze(variables_unfrozen)
            
            next_token_logits = dyna_outputs["token_logits"][:, t, n, :]
            rng, step_rng = jax.random.split(rng)
            if args.sample_argmax:
                next_token = jnp.argmax(next_token_logits, axis=-1)
            else:
                next_token = jax.random.categorical(step_rng, next_token_logits / args.temperature)
            
            token_sequence = token_sequence.at[:, t, n].set(next_token)
    
    # --- 4. Finalization (Detokenization) ---
    generated_vid = inference_genie.tokenizer.apply(
        {'params': variables['params']['tokenizer']},
        token_sequence,
        video_batch.shape[2:4],
        method=inference_genie.tokenizer.decode
[content] sample.py (python)

def _sampling_wrapper(module, batch):
    # return module.sample_maskgit(batch, args.seq_len, args.maskgit_steps, args.temperature, args.sample_argmax)
    return module.sample_causal(batch, args.seq_len, args.temperature, args.sample_argmax)
[content] sample.py (python)
    vid = video_batch[:, : args.start_frame + 1]
    # sampling_fn = jax.jit(nn.apply(_sampling_wrapper, genie)) 
    sampling_fn = nn.apply(_sampling_wrapper, inference_genie)
    rng, _rng = jax.random.split(rng)
    batch = dict(videos=vid, latent_actions=action_batch, rng=_rng)
    generated_vid = sampling_fn(
        {'params': variables_for_sampling["params"]["params"], 'cache': variables_for_sampling["cache"]},
        batch,
[content] sample.py (python)
_wrapper, genie)) 
    sampling_fn = nn.apply(_sampling_wrapper, inference_genie)
    rng, _rng = jax.random.split(rng)
    batch = dict(videos=vid, latent_actions=action_batch, rng=_rng)
    generated_vid = sampling_fn(
        {'params': variables_for_sampling["params"]["params"], 'cache': variables_for_sampling["cache"]},
        batch,
    )
    return generated_vid
[content] sample.py (python)
def _autoreg_sample(rng, video_batch, action_batch):
    """
    A clean wrapper to call the JIT-compiled sample_causal method from the Genie module.
    """
    # 1. Prepare the input batch for the model.
    # It needs the context frames, the full action sequence, and an RNG key.
    context_frames = video_batch[:, : args.start_frame + 1]
    rng, sample_rng = jax.random.split(rng)
    batch = {
        'videos': context_frames,
        'latent_actions': action_batch,
        'rng': sample_rng,
    }

    # 2. Define and JIT-compile the sampling function.
    # This is the standard and most performant way to do this in Flax.
    @jax.jit
    def sampling_fn(variables, input_batch):
        return inference_genie.apply(
            variables,
            input_batch,
            args.seq_len,
            start_frame=args.start_frame,
            temperature=args.temperature,
            sample_argmax=args.sample_argmax,
            method=Genie.sample_causal,
        )

    # 3. Call the compiled function with the correct variables and batch.
    # Because of the change in Step 1, we can pass variables_for_sampling directly.
    generated_vid = sampling_fn(variables_for_sampling, batch)
    return generated_vid
[content] genie.py (python)
def _reset_spatial_cache(cache: flax.core.FrozenDict) -> flax.core.FrozenDict:
    """Resets the spatial attention cache within a nested cache dictionary."""
    if not cache:
        return cache
    cache_unfrozen = unfreeze(cache)
    # The cache is structured like: {'STBlock_0': {'MultiHeadAttention_0': ..., 'MultiHeadAttention_1': ...}, ...}
    for block_name in cache_unfrozen:
        # The first MHA in each STBlock is spatial attention.
        if 'MultiHeadAttention_0' in cache_unfrozen[block_name]:
            spatial_attn_cache = cache_unfrozen[block_name]['MultiHeadAttention_0']
            if 'cache_index' in spatial_attn_cache:
                spatial_attn_cache['cache_index'] = jnp.array(0, dtype=jnp.int32)
            if 'cached_key' in spatial_attn_cache:
                spatial_attn_cache['cached_key'] = jnp.zeros_like(spatial_attn_cache['cached_key'])
            if 'cached_value' in spatial_attn_cache:
                spatial_attn_cache['cached_value'] = jnp.zeros_like(spatial_attn_cache['cached_value'])
    return freeze(cache_unfrozen)
[content] genie.py (python)

    def _reset_spatial_cache(cache: flax.core.FrozenDict) -> flax.core.FrozenDict:
        """Resets the spatial attention cache within a nested cache dictionary."""
        if not cache:
            return cache
        cache_unfrozen = unfreeze(cache)
        # The cache is structured like: {'STBlock_0': {'MultiHeadAttention_0': ..., 'MultiHeadAttention_1': ...}, ...}
        for block_name in cache_unfrozen:
            # The first MHA in each STBlock is spatial attention.
            if 'MultiHeadAttention_0' in cache_unfrozen[block_name]:
                spatial_attn_cache = cache_unfrozen[block_name]['MultiHeadAttention_0']
                if 'cache_index' in spatial_attn_cache:
                    spatial_attn_cache['cache_index'] = jnp.array(0, dtype=jnp.int32)
                if 'cached_key' in spatial_attn_cache:
                    spatial_attn_cache['cached_key'] = jnp.zeros_like(spatial_attn_cache['cached_key'])
                if 'cached_value' in spatial_attn_cache:
                    spatial_attn_cache['cached_value'] = jnp.zeros_like(spatial_attn_cache['cached_value'])
        return freeze(cache_unfrozen)
[content] models/dynamics.py (python)
    def init_cache(self, input_shape, dtype):
        """
        Initialize the cache for autoregressive generation.
        
        Args:
            input_shape: Shape of the input (B, seq_len, N, model_dim)
            dtype: Data type for the cache
            
        Returns:
            Cache structure for the STTransformer attention layers
        """
        B, seq_len, N, model_dim = input_shape
        head_dim = model_dim // self.num_heads
        
        # Create cache structure for each STBlock
        cache = {}
        for i in range(self.num_blocks):
            block_cache = {}
            
            # Spatial attention cache (MultiHeadAttention_0)
            # Shape: (B, seq_len, num_heads, head_dim)
            spatial_key_shape = (B, seq_len, self.num_heads, head_dim)
            spatial_value_shape = (B, seq_len, self.num_heads, head_dim)
            
            block_cache['MultiHeadAttention_0'] = {
                'cached_key': jnp.zeros(spatial_key_shape, dtype=dtype),
                'cached_value': jnp.zeros(spatial_value_shape, dtype=dtype),
                'cache_index': jnp.array(0, dtype=jnp.int32)
            }
            
            # Temporal attention cache (MultiHeadAttention_1)
            # For temporal attention, we need to swap axes: (B, N, seq_len, model_dim)
            temporal_key_shape = (B, N, seq_len, self.num_heads, head_dim)
            temporal_value_shape = (B, N, seq_len, self.num_heads, head_dim)
            
            block_cache['MultiHeadAttention_1'] = {
                'cached_key': jnp.zeros(temporal_key_shape, dtype=dtype),
                'cached_value': jnp.zeros(temporal_value_shape, dtype=dtype),
                'cache_index': jnp.array(0, dtype=jnp.int32)
            }
            
            cache[f'STBlock_{i}'] = block_cache
            
        return cache
[content] models/dynamics.py (python)
    def init_cache(self, input_shape, dtype):
        """
        Initialize the cache for autoregressive generation.
        
        Args:
            input_shape: Shape of the input (B, seq_len, N, model_dim)
            dtype: Data type for the cache
            
        Returns:
            Cache structure for the STTransformer attention layers
        """
        B, seq_len, N, model_dim = input_shape
        head_dim = model_dim // self.num_heads
        
        # Create cache structure for each STBlock
        cache = {}
        for i in range(self.num_blocks):
            block_cache = {}
            
            # Spatial attention cache (MultiHeadAttention_0)
            # Shape: (B, seq_len, num_heads, head_dim)
            spatial_key_shape = (B, seq_len, self.num_heads, head_dim)
            spatial_value_shape = (B, seq_len, self.num_heads, head_dim)
            
            block_cache['MultiHeadAttention_0'] = {
                'cached_key': jnp.zeros(spatial_key_shape, dtype=dtype),
                'cached_value': jnp.zeros(spatial_value_shape, dtype=dtype),
                'cache_index': jnp.array(0, dtype=jnp.int32)
            }
            
            # Temporal attention cache (MultiHeadAttention_1)
            # For temporal attention, input is swapped: (B, N, seq_len, model_dim)
            # So cache should be: (B, N, seq_len, num_heads, head_dim)
            temporal_key_shape = (B, N, seq_len, self.num_heads, head_dim)
            temporal_value_shape = (B, N, seq_len, self.num_heads, head_dim)
            
            block_cache['MultiHeadAttention_1'] = {
                'cached_key': jnp.zeros(temporal_key_shape, dtype=dtype),
                'cached_value': jnp.zeros(temporal_value_shape, dtype=dtype),
                'cache_index': jnp.array(0, dtype=jnp.int32)
            }
            
            cache[f'STBlock_{i}'] = block_cache
            
        return cache
[content] models/dynamics.py (python)
            # Input: (B, N, seq_len, model_dim) -> treated as (B*N, seq_len, model_dim)
            # Cache: (B*N, seq_len, num_heads, head_dim)
            temporal_key_shape = (B * N, seq_len, self.num_heads, head_dim)
            temporal_value_shape = (B * N, seq_len, self.num_heads, head_dim)
[content] models/dynamics.py (python)
            # Input: (B, seq_len, N, model_dim) -> treated as (B*N, seq_len, model_dim)
            # Cache: (B*N, seq_len, num_heads, head_dim)
            spatial_key_shape = (B * N, seq_len, self.num_heads, head_dim)
            spatial_value_shape = (B * N, seq_len, self.num_heads, head_dim)
[content] models/dynamics.py (python)
            # For temporal attention, input is swapped: (B, N, seq_len, model_dim)
            # So cache should be: (B, N, seq_len, num_heads, head_dim)
            temporal_key_shape = (B, N, seq_len, self.num_heads, head_dim)
            temporal_value_shape = (B, N, seq_len, self.num_heads, head_dim)
[content] models/dynamics.py (python)
    def init_cache(self, input_shape, dtype):
        """
        Initialize the cache for autoregressive generation.
        
        Args:
            input_shape: Shape of the input (B, seq_len, N, model_dim)
            dtype: Data type for the cache
            
        Returns:
            Cache structure for the STTransformer attention layers
        """
        B, seq_len, N, model_dim = input_shape
        head_dim = model_dim // self.num_heads
        
        # Create cache structure for each STBlock
        cache = {}
        for i in range(self.num_blocks):
            block_cache = {}
            
            # Spatial attention cache (MultiHeadAttention_0)
            # Input: (B, seq_len, N, model_dim) -> treated as (B*N, seq_len, model_dim)
            # Cache: (B*N, seq_len, num_heads, head_dim)
            spatial_key_shape = (B * N, seq_len, self.num_heads, head_dim)
            spatial_value_shape = (B * N, seq_len, self.num_heads, head_dim)
            
            block_cache['MultiHeadAttention_0'] = {
                'cached_key': jnp.zeros(spatial_key_shape, dtype=dtype),
                'cached_value': jnp.zeros(spatial_value_shape, dtype=dtype),
                'cache_index': jnp.array(0, dtype=jnp.int32)
            }
            
            # Temporal attention cache (MultiHeadAttention_1)
            # Input: (B, N, seq_len, model_dim) -> treated as (B*N, seq_len, model_dim)
            # Cache: (B*N, seq_len, num_heads, head_dim)
            temporal_key_shape = (B * N, seq_len, self.num_heads, head_dim)
            temporal_value_shape = (B * N, seq_len, self.num_heads, head_dim)
            
            block_cache['MultiHeadAttention_1'] = {
                'cached_key': jnp.zeros(temporal_key_shape, dtype=dtype),
                'cached_value': jnp.zeros(temporal_value_shape, dtype=dtype),
                'cache_index': jnp.array(0, dtype=jnp.int32)
            }
            
            cache[f'STBlock_{i}'] = block_cache
            
        return cache
[content] models/dynamics.py (python)
            # Temporal attention cache (MultiHeadAttention_1)
            # Input: (B, N, seq_len, model_dim) -> treated as (B*N, seq_len, model_dim)
            # Cache: (B*N, seq_len, num_heads, head_dim)
            temporal_key_shape = (B * N, seq_len, self.num_heads, head_dim)
            temporal_value_shape = (B * N, seq_len, self.num_heads, head_dim)
            
            block_cache['MultiHeadAttention_1'] = {
                'cached_key': jnp.zeros(temporal_key_shape, dtype=dtype),
                'cached_value': jnp.zeros(temporal_value_shape, dtype=dtype),
                'cache_index': jnp.array(0, dtype=jnp.int32)
            }
[content] models/dynamics.py (python)
    def init_cache(self, input_shape, dtype):
        """
        Initialize the cache for autoregressive generation.
        Only creates spatial attention caches (MultiHeadAttention_0).
        
        Args:
            input_shape: Shape of the input (B, seq_len, N, model_dim)
            dtype: Data type for the cache
            
        Returns:
            Cache structure for the STTransformer attention layers
        """
        B, seq_len, N, model_dim = input_shape
        head_dim = model_dim // self.num_heads
        
        # Create cache structure for each STBlock
        cache = {}
        for i in range(self.num_blocks):
            block_cache = {}
            
            # Spatial attention cache (MultiHeadAttention_0) only
            # Input: (B, seq_len, N, model_dim) -> treated as (B*N, seq_len, model_dim)
            # Cache: (B*N, seq_len, num_heads, head_dim)
            spatial_key_shape = (B * N, seq_len, self.num_heads, head_dim)
            spatial_value_shape = (B * N, seq_len, self.num_heads, head_dim)
            
            block_cache['MultiHeadAttention_0'] = {
                'cached_key': jnp.zeros(spatial_key_shape, dtype=dtype),
                'cached_value': jnp.zeros(spatial_value_shape, dtype=dtype),
                'cache_index': jnp.array(0, dtype=jnp.int32)
            }
            
            # No temporal attention cache (MultiHeadAttention_1) - not needed for spatial-only caching
            
            cache[f'STBlock_{i}'] = block_cache
            
        return cache
[content] genie.py (python)
                dyna_outputs, mutable_vars = self.dynamics.apply({'params': self.dynamics.params, 'cache': cache}, dyna_inputs, training=False, mutable=['cache'])
                cache = mutable_vars['cache']
[content] genie.py (python)
    # Convert FrozenDict to dict for updating
    tokenizer_params_dict = unfreeze(train_state.params["params"]["tokenizer"])
    tokenizer_params_dict.update(restored_tokenizer_params)
    train_state.params["params"]["tokenizer"] = freeze(tokenizer_params_dict)
[content] genie.py (python)
        # Convert FrozenDict to dict for updating
        lam_params_dict = unfreeze(train_state.params["params"]["lam"])
        lam_params_dict.update(restored_lam_params)
        train_state.params["params"]["lam"] = freeze(lam_params_dict)