def METHOD_NAME(layer, fn):
  """Traverses the layer tree and invokes fn(node) on each node.
  Args:
    layer: a BaseLayer.
    fn: a function of (layer, layer_theta) -> None.
  """
  if isinstance(layer, (list, tuple)):
    for layer_i in layer:
      METHOD_NAME(layer_i, fn)
    return
  if isinstance(layer, (dict, py_utils.NestedMap)):
    for _, layer_i in layer.items():
      METHOD_NAME(layer_i, fn)
    return
  with tf.name_scope(layer.params.name):
    fn(layer)
    # Traverse all children in alphabetical order.
    for _, child in sorted(layer.children.items()):
      METHOD_NAME(child, fn)