From e796ff4a7b2a5b2b5433a685263441dcfde0226d Mon Sep 17 00:00:00 2001
From: Sylvester Kaczmarek
 <16242628+sylvesterkaczmarek@users.noreply.github.com>
Date: Sat, 15 Aug 2026 17:00:15 +0100
Subject: [PATCH] Allow hashable axis names in shape checks

Penzai named arrays support arbitrary hashable axis names, including TmpPosAxisMarker for temporary positional-axis bindings. check_structure still asserted that named-axis keys captured through an unpacked **var(...) binding were strings, both when inlining a solved mapping and when re-associating the final match result.

Remove those string-only assertions while preserving the existing matching logic.

Regression coverage includes repeated unpacked named-axis variables with a TmpPosAxisMarker and the Linear reproduction from issue #132.

Closes google-deepmind/penzai#132
---
 penzai/core/shapecheck.py                     |  2 -
 .../shapecheck_hashable_axis_names_test.py    | 59 +++++++++++++++++++
 2 files changed, 59 insertions(+), 2 deletions(-)
 create mode 100644 tests/core/shapecheck_hashable_axis_names_test.py

diff --git a/penzai/core/shapecheck.py b/penzai/core/shapecheck.py
index 78580ea..dbd3004 100644
--- a/penzai/core/shapecheck.py
+++ b/penzai/core/shapecheck.py
@@ -618,7 +618,6 @@ def _named_inline_multidimvars(
     binding = solutions[key.name]
     assert isinstance(binding.value, dict)
     for subkey, subval in binding.value.items():
-      assert isinstance(subkey, str)
       if subkey in new_pattern:
         return (
             _UnsatisfiedConstraint(
@@ -1038,7 +1037,6 @@ def add_constraints(keypath, pattern: Any, value: Any):
           ):
             found = solutions[name[0]].value[name[1]]
         else:
-          assert isinstance(name[1], str)
           if isinstance(solutions[name[0]].value, dict):
             found = solutions[name[0]].value.get(name[1])
         if found != binding.value:
diff --git a/tests/core/shapecheck_hashable_axis_names_test.py b/tests/core/shapecheck_hashable_axis_names_test.py
new file mode 100644
index 0000000..0bbc814
--- /dev/null
+++ b/tests/core/shapecheck_hashable_axis_names_test.py
@@ -0,0 +1,59 @@
+# Copyright 2026 The Penzai Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Regression tests for hashable named axes in shape checking."""
+
+from absl.testing import absltest
+import jax
+from penzai import pz
+
+
+class ShapecheckHashableAxisNamesTest(absltest.TestCase):
+
+  def test_repeated_named_unpack_accepts_hashable_axis_name(self):
+    temp_axis = pz.nx.TmpPosAxisMarker()
+
+    match = pz.chk.check_structure(
+        value={
+            "a": pz.chk.ArraySpec(named_shape={"feature": 2, temp_axis: 3}),
+            "b": pz.chk.ArraySpec(named_shape={"feature": 4, temp_axis: 3}),
+        },
+        pattern={
+            "a": pz.chk.ArraySpec(
+                named_shape={"feature": 2, **pz.chk.var("batch_axes")}
+            ),
+            "b": pz.chk.ArraySpec(
+                named_shape={"feature": 4, **pz.chk.var("batch_axes")}
+            ),
+        },
+    )
+
+    self.assertEqual(match["batch_axes"], {temp_axis: 3})
+
+  def test_linear_accepts_temporary_positional_axis_marker(self):
+    temp_axis = pz.nx.TmpPosAxisMarker()
+    layer = pz.nn.Linear.from_config(
+        "my_layer",
+        jax.random.key(0),
+        input_axes={"in_axis": 2},
+        output_axes={"out_axis": 3},
+    )
+
+    result = layer(pz.nx.zeros({"in_axis": 2, temp_axis: 3}))
+
+    self.assertEqual(result.named_shape, {temp_axis: 3, "out_axis": 3})
+
+
+if __name__ == "__main__":
+  absltest.main()