improve(asserts): clarify bare Ellipsis usage in assert_shape error

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):

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)