def initialize_model(*, model_dict: Dict[str, nn.Module],
                     input_spec: Sequence[Union[Tuple[Tuple[int, ...],
                                                      jnp.dtype],
                                                Tuple[int, ...], None]],
                     config: ml_collections.ConfigDict, rng: jnp.ndarray,
                     is_train: bool):
  """Initializes parameters and model state.

  Args:
    model_dict: Dictionary of the models.
    input_spec: An iterable of (shape, dtype) pairs specifying the shape and
      dtype of the inputs. If unspecified the dtype is float32.
    config: Configurations of the initialization.
    rng: Jax rng keys.
    is_train: whether is in training mode.

  Returns:
    Initial params, Init model_state, and number of trainable_params.
  """
  batch_size = config.batch_size
  if not is_train:
    batch_size = config.get('eval_batch_size', batch_size)
  dev_batch_size = batch_size // jax.device_count()
  dummy_input = []
  for spec in input_spec:
    if spec is not None:
      in_st = debug_utils.input_spec_to_jax_shape_dtype_struct(
          spec, batch_size=dev_batch_size)
      dummy_input.append(jnp.zeros(in_st.shape, in_st.dtype))
    else:
      dummy_input.append(None)
  # 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."""
    d_rng, g_rng = jax.random.split(rng, 2)
    generator_variables = model_dict['generator'].init(g_rng, *dummy_input)
    g_model_state, g_params = generator_variables.pop('params')

    d_dummy_input = jax.tree_util.tree_map(
        lambda x, y: jnp.concatenate([x, y], axis=0), dummy_input, dummy_input)
    discriminator_variables = model_dict['discriminator'].init(
        d_rng, *d_dummy_input)
    d_model_state, d_params = discriminator_variables.pop('params')

    state_dict = dict(
        generator=g_model_state,
        discriminator=d_model_state)
    init_params_dict = dict(
        generator=g_params, discriminator=d_params)

    return state_dict, init_params_dict

  state_dict, init_params_dict = _initialize_model(rng)
  return state_dict, init_params_dict
Initializes parameters and model state.

  Args:
    model_dict: Dictionary of the models.
    input_spec: An iterable of (shape, dtype) pairs specifying the shape and
      dtype of the inputs. If unspecified the dtype is float32.
    config: Configurations of the initialization.
    rng: Jax rng keys.
    is_train: whether is in training mode.

  Returns:
    Initial params, Init model_state, and number of trainable_params.
  
videogvt/trainers/vqgan_trainer.py