def _sizes(x):
  """Get a structure of sizes for a structure of nested arrays."""
  def size(x):
    try:
      return x.size
    except Exception:  # pylint: disable=broad-except
      return 0
  return fastmath.nested_map(size, x)
Get a structure of sizes for a structure of nested arrays.
trax/supervised/trainer_lib.py