def assert_shape_equals(array, shape):
  """Asserts that an array has the given shape."""
  assert array.shape == shape, (
      'Invalid shape {}; expected {}.'.format(array.shape, shape)
  )
Asserts that an array has the given shape.
trax/shapes.py