def generate_model_masks(
    depth,
    mask = None,
    masked_layer_indices = None):
  """Creates empty masks for this model, or initializes with existing mask.

  Args:
    depth: Number of layers in the model.
    mask: Existing model mask for layers in this model, if not given, all
      module masks are initialized to None.
    masked_layer_indices: The layer indices of layers in model to be masked, or
      all if None.

  Returns:
    A model mask, with None where no mask is given for a model layer, or that
    specific layer is indicated as not to be masked by the masked_layer_indices
    parameter.
  """
  if depth <= 0:
    raise ValueError(f'Invalid model depth: {depth}')

  if mask is None:
    mask = {f'MaskedModule_{i}': None for i in range(depth)}

  # Have to explicitly check for None to differentiate from empty array.
  if masked_layer_indices is not None:
    # Check none of the indices are outside of model's layer bounds.
    if any(i < 0 or i >= depth for i in masked_layer_indices):
      raise ValueError(
          f'Invalid indices for given depth ({depth}): {masked_layer_indices}')
    mask = {
        f'MaskedModule_{i}': mask[f'MaskedModule_{i}']
        for i in masked_layer_indices
    }

  return mask