diff --git a/penzai/core/shapecheck.py b/penzai/core/shapecheck.py index 78580ea..a996fa0 100644 --- a/penzai/core/shapecheck.py +++ b/penzai/core/shapecheck.py @@ -93,7 +93,7 @@ class DimVar(Mapping): and the inner name is the named shape of a single axis in that collection. """ - name: str | tuple[str, str | int] + name: str | tuple[str, named_axes.AxisName] def __len__(self) -> int: return 1 @@ -148,7 +148,7 @@ class KnownDim: from_keypath: Optional keypath that indicates where this size was bound. """ - name: str | tuple[str, str | int] + name: str | tuple[str, named_axes.AxisName] size: int from_keypath: str | None = None @@ -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( @@ -972,9 +971,13 @@ def add_constraints(keypath, pattern: Any, value: Any): ) elif remaining_value_keys: # No unpack pattern, but we ran out of keys. + # Use list() instead of sorted() because axis names may be + # non-string hashables (e.g. TmpPosAxisMarker) that are not + # mutually orderable. OrderedSet preserves insertion order, which + # gives a deterministic message without requiring comparison. subfailures.append( " Unexpected names in value's named shape:" - f" {sorted(remaining_value_keys)}" + f" {list(remaining_value_keys)}" ) if subfailures: failures.append(( @@ -1038,7 +1041,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: @@ -1090,7 +1092,7 @@ class DimensionVariableSubstitution: axis name to sizes. """ - size_variables: dict[str | tuple[str, str | int], int | DimVar] + size_variables: dict[str | tuple[str, named_axes.AxisName], int | DimVar] sequence_variables: dict[str, tuple[int | DimVar | MultiDimVar, ...]] mapping_variables: dict[ str, dict[Any | MultiDimVar, int | DimVar | RemainingAxisPlaceholder] diff --git a/tests/core/shapecheck_test.py b/tests/core/shapecheck_test.py index 80afc41..cee8ac1 100644 --- a/tests/core/shapecheck_test.py +++ b/tests/core/shapecheck_test.py @@ -748,6 +748,104 @@ def test_get_and_substitute_dimension_variables(self): } self.assertEqual(res, expected) + def test_nonstring_axis_names_in_check_structure(self): + """Regression test for issue #132. + + pz.chk.check_structure should work when named array axes use non-string + AxisNames (e.g. TmpPosAxisMarker), since AxisName = Hashable allows any + hashable as an axis name. + + The crash occurs when a solved **var() MultiDimVar binding contains + non-string keys and is passed back via known_vars for inlining. + """ + nonstr_axis = pz.nx.TmpPosAxisMarker() + + # Step 1: solve the batch var against an array with a non-string axis name. + input_arr = pz.nx.zeros({"in_axis": 2, nonstr_axis: 3}) + input_pattern = pz.chk.ArraySpec( + named_shape={"in_axis": 2, **pz.chk.var("batch")} + ) + dimvars = pz.chk.check_structure(input_arr, input_pattern) + self.assertEqual(dict(dimvars), {"batch": {nonstr_axis: 3}}) + + # Step 2: use those dimvars as known_vars for a second check (this + # previously crashed at shapecheck.py:621 with AssertionError). + output_arr = pz.nx.zeros({nonstr_axis: 3, "out_axis": 5}) + output_pattern = pz.chk.ArraySpec( + named_shape={**pz.chk.var("batch"), "out_axis": 5} + ) + result = pz.chk.check_structure( + output_arr, output_pattern, known_vars=dimvars + ) + # The result includes the known_vars variables re-confirmed as consistent. + self.assertEqual(dict(result), {"batch": {nonstr_axis: 3}}) + + def test_linear_layer_with_tmpposaxismarker_batch(self): + """Regression test for issue #132: Linear layer with TmpPosAxisMarker batch axis. + + This is the exact repro from the issue report: calling a Linear layer on a + NamedArray that has a TmpPosAxisMarker as one of its named axes. + """ + my_layer = pz.nn.Linear.from_config( + "my_layer", + jax.random.key(0), + input_axes={"in_axis": 2}, + output_axes={"out_axis": 3}, + ) + nonstr_batch_axis = pz.nx.TmpPosAxisMarker() + # This previously crashed with AssertionError at shapecheck.py:621. + result = my_layer(pz.nx.zeros({"in_axis": 2, nonstr_batch_axis: 3})) + self.assertIn(nonstr_batch_axis, result.named_axes) + self.assertIn("out_axis", result.named_axes) + + def test_nonstring_axis_name_mismatch_error_message(self): + """Non-string axis names in error paths should produce StructureMismatchError. + + When a named shape mismatch occurs and the value contains a single + non-string AxisName not present in the pattern (the 'unexpected names' + error path), check_structure must raise StructureMismatchError rather than + crashing with TypeError inside sorted(). This exercises the error-message + path in the named constraint solver (shapecheck.py line ~976). + """ + nonstr_axis = pz.nx.TmpPosAxisMarker() + + # Array has an extra non-string axis that the pattern does not mention. + arr = pz.nx.zeros({nonstr_axis: 3, "out_axis": 5}) + pattern = pz.chk.ArraySpec(named_shape={"out_axis": 5}) + with self.assertRaises(pz.chk.StructureMismatchError): + pz.chk.check_structure(arr, pattern) + + # Non-string axis present with the wrong size. + arr2 = pz.nx.zeros({nonstr_axis: 3, "out_axis": 5}) + pattern2 = pz.chk.ArraySpec(named_shape={nonstr_axis: 99, "out_axis": 5}) + with self.assertRaises(pz.chk.StructureMismatchError): + pz.chk.check_structure(arr2, pattern2) + + def test_multiple_nonstring_axis_names_unexpected_error_path(self): + """Two non-orderable non-string axis names in unexpected-names error path. + + When the value has extra non-string axis names not present in the pattern + and no unpack var(**) to absorb them, the error-formatting path at + shapecheck.py line ~976 previously called sorted(remaining_value_keys), + which crashes with TypeError for non-orderable types (e.g. two distinct + TmpPosAxisMarker instances). The fix replaces sorted() with list() since + OrderedSet already has deterministic insertion order. + + This is in-scope for issue #132: the crash was latent on main but becomes + reachable in realistic use once non-string axis names can flow through + check_structure (which is what #132 enables). + """ + axis_a = pz.nx.TmpPosAxisMarker() + axis_b = pz.nx.TmpPosAxisMarker() + + # Two non-orderable non-string axis names that are unexpected by the pattern. + # Previously crashed with TypeError inside sorted(); now must raise + # StructureMismatchError with a readable message. + arr = pz.nx.zeros({axis_a: 3, axis_b: 5, "out": 7}) + pattern = pz.chk.ArraySpec(named_shape={"out": 7}) + with self.assertRaises(pz.chk.StructureMismatchError): + pz.chk.check_structure(arr, pattern) + if __name__ == "__main__": absltest.main()