def __init__(self, shape, dtype=np.float32):
    """Creates a `ShapeDtype` instance, with canonicalized `shape` and `dtype`.

    Args:
      shape: A tuple or list, each element of which is an int or, less often,
          `None`.
      dtype: A `dtype` object, either from NumPy or TensorFlow.

    Returns:
      A `ShapeDtype` instance whose `shape` is a tuple and `dtype` is a NumPy
      `dtype` object.
    """
    # Canonicalize shape and dtype.
    if isinstance(shape, tf.TensorShape):
      shape = shape.as_list()
    if isinstance(shape, list):
      shape = tuple(shape)
    if not isinstance(shape, tuple):
      raise TypeError('shape must be tuple or list; got: {}'.format(shape))
    if isinstance(dtype, tf.DType):
      dtype = dtype.as_numpy_dtype

    self.shape = shape
    self.dtype = dtype
Creates a `ShapeDtype` instance, with canonicalized `shape` and `dtype`.

    Args:
      shape: A tuple or list, each element of which is an int or, less often,
          `None`.
      dtype: A `dtype` object, either from NumPy or TensorFlow.

    Returns:
      A `ShapeDtype` instance whose `shape` is a tuple and `dtype` is a NumPy
      `dtype` object.
    
trax/shapes.py