def _tree_broadcast(to, s):
  """Broadcasts `s` to the nested structure `to`."""
  if not isinstance(to, (list, tuple, dict)):
    if not isinstance(s, (int, type(None))):
      raise ValueError
    return s
  if isinstance(s, (int, type(None))):
    return tf.nest.map_structure(lambda x: s, to)
  if isinstance(to, (list, tuple)):
    if len(to) != len(s):
      raise ValueError
    new_s = [_tree_broadcast(x, y) for x, y in zip(to, s)]
    if isinstance(to, tuple):
      new_s = tuple(new_s)
    return new_s
  elif isinstance(to, dict):
    return {k: _tree_broadcast(to[k], s[k]) for k in to}
  else:
    raise TypeError("Unsupported type %s" % type(to))
Broadcasts `s` to the nested structure `to`.
trax/tf_numpy/extensions/extensions.py