def _is_at_level(obj, level):
  """Checks if `obj` is an at level `level`."""
  is_leaf = not isinstance(obj, (list, tuple, dict))
  if level == 0 or is_leaf:
    return (level == 0) == is_leaf

  if isinstance(obj, dict):
    elems = obj.values()
  else:
    elems = obj
  return elems and all(_is_at_level(x, level - 1) for x in elems)
Checks if `obj` is an at level `level`.
trax/fastmath/numpy.py