From d4460fe6bb1d9b63669e9814ea7c583ff4ac35de Mon Sep 17 00:00:00 2001
From: abhinav <244986440+erensh27@users.noreply.github.com>
Date: Sun, 9 Aug 2026 13:07:36 +0530
Subject: [PATCH] fix(asserts): clarify bare Ellipsis usage in assert_shape
 error

Passing  directly to assert_shape raised
'expected shapes should be a list or tuple of ints, got Ellipsis',
which implies a type mistake even though  is a valid wildcard
shape. Special-case Ellipsis in the error message and point users at
the  form.
---
 chex/_src/asserts.py      | 5 ++++-
 chex/_src/asserts_test.py | 9 +++++++++
 2 files changed, 13 insertions(+), 1 deletion(-)

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)