Helper: sum a list of arrays or nested arrays.

def _nested_op(inputs, op):  # pylint: disable=invalid-name
  """Helper: sum a list of arrays or nested arrays."""
  # First the simple non-nested case.
  if not isinstance(inputs[0], (list, tuple)):
    return op(inputs)
  # In the nested case, sum on each axis separately.
  result_list = []
  for i in range(len(inputs[0])):
    result_list.append(_nested_op([x[i] for x in inputs], op=op))
  if isinstance(inputs[0], list):
    return result_list
  return tuple(result_list)