[content] sample.py (python)
def _autoreg_sample(rng, video_batch, action_batch):
    # --- 1. Define the JIT-compiled step function ---
    @jax.jit
    def predict_token_fn(variables, inputs):
        """
        This function calls the dynamics submodule's apply method.
        """
        # The `method` argument can be a direct reference to a submodule's __call__
        return inference_genie.dynamics.apply(
            variables,
            inputs,
            training=False,
            mutable=['cache']
        )

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

    # Tokenize by applying the tokenizer submodule directly
    tokenizer_out = inference_genie.tokenizer.apply(
        {'params': variables['params']['tokenizer']},
        context_frames,
        training=False,
        method=inference_genie.tokenizer.vq_encode
    )
    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)

    # Get action codes by applying the LAM's vector quantizer submodule
    action_codes = inference_genie.lam.vq.apply(
        {'params': variables['params']['lam']['vq']},
        action_batch,
        method=inference_genie.lam.vq.get_codes
    )

    # --- 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,
            }
            
            # Here we need to pass the variables relevant to the dynamics model
            dynamics_variables = {'params': variables['params']['dynamics'], 'cache': variables['cache']['dynamics']}
            
            # Call the pure, JIT-compiled function
            dyna_outputs, updated_dynamics_cache = predict_token_fn(dynamics_variables, step_inputs)
            
            # Manually update the main variables dictionary with the new cache state
            variables['cache']['dynamics'] = updated_dynamics_cache['cache']
            
            # Extract logits and sample
            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) ---
    # Detokenize by applying the tokenizer submodule directly
    generated_vid = inference_genie.tokenizer.apply(
        {'params': variables['params']['tokenizer']},
        token_sequence,
        video_batch.shape[2:4], # Pass target HW
        method=inference_genie.tokenizer.decode
    )
    return generated_vid