def _dims_of_shape(shape):
  """Converts `shape` to a tuple of dimensions."""
  if type(shape) in (list, tuple):
    return shape
  elif isinstance(shape, ScalarShape):
    return ()
  else:
    raise TypeError(type(shape))
Converts `shape` to a tuple of dimensions.
trax/tf_numpy/jax_tests/test_util.py