From fe7d9c0cf799f5030a4f64684dd2a7df26945cbb Mon Sep 17 00:00:00 2001
From: Dmitrii Kochkov <dkochkov@google.com>
Date: Fri, 15 May 2026 14:00:06 -0700
Subject: [PATCH] Fixed validate_coordinate check on Field to skip DummyAxis
 not only when it is provided explicitly, but also when it is a component of a
 CartesianProduct.

PiperOrigin-RevId: 916171079
---
 coordax/fields.py      |  4 ++--
 coordax/fields_test.py | 31 ++++++++++++++++++++++++++-----
 2 files changed, 28 insertions(+), 7 deletions(-)

diff --git a/coordax/fields.py b/coordax/fields.py
index 65afff3..c9a9325 100644
--- a/coordax/fields.py
+++ b/coordax/fields.py
@@ -691,8 +691,8 @@ def _validate_matching_coords(
     """Validate that given coordinates are all found on this field."""
     axes = []
     for part in dims_or_coords:
-      if isinstance(part, Coordinate) and not isinstance(part, DummyAxis):
-        axes.extend(part.axes)
+      if isinstance(part, Coordinate):
+        axes.extend([c for c in part.axes if not isinstance(c, DummyAxis)])
 
     for c in axes:
       [dim] = c.dims
diff --git a/coordax/fields_test.py b/coordax/fields_test.py
index b9eed2a..7764bf8 100644
--- a/coordax/fields_test.py
+++ b/coordax/fields_test.py
@@ -283,6 +283,30 @@ def test_field_repr(self):
           full_unwrap=False,
           should_raise_on_untag=False,
       ),
+      dict(
+          testcase_name='composed_dummy_axis',
+          array=np.arange(4 * 5).reshape(4, 5),
+          tags=('x', 'y'),
+          untags=(
+              coordax.coords.compose(
+                  coordax.DummyAxis('x', 4), coordax.DummyAxis('y', 5)
+              ),
+          ),
+          full_unwrap=True,
+          should_raise_on_untag=False,
+      ),
+      dict(
+          testcase_name='composed_dummy_axis_missing',
+          array=np.arange(4 * 5).reshape(4, 5),
+          tags=('x', 'y'),
+          untags=(
+              coordax.coords.compose(
+                  coordax.DummyAxis('x', 4), coordax.DummyAxis('missing', 5)
+              ),
+          ),
+          full_unwrap=False,
+          should_raise_on_untag=True,
+      ),
   )
   def test_tag_then_untag_by(
       self,
@@ -731,9 +755,7 @@ def test_get_coordinate(self):
     )
     with self.subTest('default'):
       actual = coordax.get_coordinate(field)
-      expected = coordax.coords.compose(
-          *[axes[d] for d in ['x', 'y', 'z']]
-      )
+      expected = coordax.coords.compose(*[axes[d] for d in ['x', 'y', 'z']])
       self.assertEqual(actual, expected)
 
     with self.subTest('with_positional_dims'):
@@ -952,8 +974,7 @@ def test_sel_raises_on_unused_indexer(self):
     y = coordax.LabeledAxis('y', np.array([100, 200, 300]))
     field = coordax.field(jnp.zeros((2, 3)), x, y)
     with self.assertRaisesRegex(
-        ValueError,
-        re.escape("Indexers {'z'} were not processed")
+        ValueError, re.escape("Indexers {'z'} were not processed")
     ):
       field.sel(z=slice(0, 20), x=10)