def flattened_traversal(fn):
  """Utility function to create masks for optax.

  Adapted from:
  https://github.com/google/flax/blob/main/docs/flip/1009-optimizer-api.md

  Args:
    fn: a filter function e.g., lambda path, _: path[-1] == 'bias'

  Returns:
    The function to produce the mask.
  """
  def mask(data):
    flat = flax.traverse_util.flatten_dict(data)
    return flax.traverse_util.unflatten_dict(
        {k: fn(k, v) for k, v in flat.items()})

  return mask
Utility function to create masks for optax.

  Adapted from:
  https://github.com/google/flax/blob/main/docs/flip/1009-optimizer-api.md

  Args:
    fn: a filter function e.g., lambda path, _: path[-1] == 'bias'

  Returns:
    The function to produce the mask.
  
videogvt/train_lib/train_utils.py