def inflate_2d_to_3d(tree_2d, tree_3d, *, mode='central'):
  """Inflate params of 2D CNN to 3D CNN.

  Note: If tree_2d is not a frozen dict and tree_3d is a regular dict,
  tree_2d is unfrozen before inflation.

  Args:
    tree_2d: The tree of 2D params.
    tree_3d: The tree of 3D params.
    mode: One of 'average' or 'central', the mode to use for parameter
      inflation.

  Returns:
    The inflated 3D tree.
  """
  assert mode in ['average', 'central'], f'Unknown inflation mode {mode}'

  if isinstance(tree_2d, flax.core.FrozenDict) and isinstance(tree_3d, dict):
    tree_2d = flax.core.unfreeze(tree_2d)

  def inflate_param(param_2d, param_3d):
    if param_2d.shape == param_3d.shape:
      return param_2d
    if param_2d.shape == param_3d.shape[1:]:
      if mode == 'average':
        param_2d_replicated = jnp.broadcast_to(param_2d[None], param_3d.shape)
        return param_2d_replicated / param_3d.shape[0]
      elif mode == 'central':
        param_3d_zero = jnp.zeros_like(param_3d)
        mid = param_3d.shape[0] // 2
        return param_3d_zero.at[mid].set(param_2d)
    logging.warn('Unable to inflate from %s to %s, left as is.', param_2d.shape,
                 param_3d.shape)
    return param_3d

  return jax.tree_util.tree_map(inflate_param, tree_2d, tree_3d)
Inflate params of 2D CNN to 3D CNN.

  Note: If tree_2d is not a frozen dict and tree_3d is a regular dict,
  tree_2d is unfrozen before inflation.

  Args:
    tree_2d: The tree of 2D params.
    tree_3d: The tree of 3D params.
    mode: One of 'average' or 'central', the mode to use for parameter
      inflation.

  Returns:
    The inflated 3D tree.
  
videogvt/train_lib/vqgan_train_lib.py