def _shapes(x):
  """Gets a structure of shapes for a structure of nested arrays."""
  def shape(x):
    try:
      return tuple([int(i) for i in x.shape])
    except Exception:  # pylint: disable=broad-except
      return ()
  return tuple(nested_map(shape, x))
Gets a structure of shapes for a structure of nested arrays.
trax/layers/base.py