diff --git a/chex/_src/asserts.py b/chex/_src/asserts.py index 0e114e0..81dfbe2 100644 --- a/chex/_src/asserts.py +++ b/chex/_src/asserts.py @@ -632,9 +632,12 @@ def assert_shape( not match ``expected_shapes``. """ if not isinstance(expected_shapes, (list, tuple)): + suffix = "" + if expected_shapes is Ellipsis: + suffix = " For a wildcard shape, pass a tuple, e.g. (...,) instead of a bare Ellipsis (...)." raise AssertionError( "Error in shape compatibility check: expected shapes should be a list " - f"or tuple of ints, got {expected_shapes}.") + f"or tuple of ints, got {expected_shapes}.{suffix}") # Ensure inputs and expected shapes are sequences. if not isinstance(inputs, collections.abc.Sequence): diff --git a/chex/_src/asserts_test.py b/chex/_src/asserts_test.py index c6f6d20..85746f3 100644 --- a/chex/_src/asserts_test.py +++ b/chex/_src/asserts_test.py @@ -547,6 +547,15 @@ def test_multiple_ellipses(self, array, expected_shape): '`expected_shape` may not contain more than one ellipsis, but got .+'): asserts.assert_shape(array, expected_shape) + def test_bare_ellipsis_message_suggests_tuple_form(self): + array = array_from_shape(2, 3) + with self.assertRaisesRegex( + AssertionError, + 'expected shapes should be a list or tuple of ints, got Ellipsis. ' + r'For a wildcard shape, pass a tuple, e.g. \(\.\.\.,\) instead of a ' + 'bare Ellipsis \\(...\\)'): + asserts.assert_shape(array, Ellipsis) + def rank_array(n): return np.zeros(shape=[2] * n)