def initialize_model(*,
                     model: nn.Module,
                     dummy_inputs: jnp.ndarray,
                     rng: jnp.ndarray):
  """Initializes parameters and model state.

  Args:
    model: The model.
    dummy_inputs: Dummy model inputs.
    rng: Jax rng keys.

  Returns:
    Init model_state, Initial params.
  """
  # We want all parameters to be created in host RAM, not on any device, they'll
  # be sent there later as needed, otherwise we already encountered two
  # situations where we allocate them twice.
  @functools.partial(jax.jit, backend='cpu')
  def _initialize_model(rng):
    """Initialization function to be jitted."""
    model_variables = model.init(rng, dummy_inputs)
    model_state = dict(model_variables)
    init_params = model_state.pop('params')
    return model_state, init_params
  model_state, init_params = _initialize_model(rng)
  return model_state, init_params
Initializes parameters and model state.

  Args:
    model: The model.
    dummy_inputs: Dummy model inputs.
    rng: Jax rng keys.

  Returns:
    Init model_state, Initial params.
  
videogvt/trainers/lmgvt_trainer.py