def initialize_cache(inputs: np.ndarray,
                     *,
                     task: Task,
                     model: nn.Module,
                     input_spec: InputSpec,
                     tokenizer_fn: TokenizerFn,
                     config: ml_collections.ConfigDict):
  """Initializes the cache variables for self-attention layers."""
  dummy_inputs = _get_dummy_inputs(input_spec, task, tokenizer_fn, config,
                                   batch_size=inputs.shape[0])
  initial_variables = model.init(jax.random.PRNGKey(0), dummy_inputs)
  return initial_variables['cache']
Initializes the cache variables for self-attention layers.
videogvt/trainers/lmgvt_trainer.py