Skip to content

fix: allow non-string AxisNames in check_structure (issue #132)#138

Open
Nas01010101 wants to merge 1 commit into
google-deepmind:mainfrom
Nas01010101:fix/issue-132-nonstr-axis-names
Open

fix: allow non-string AxisNames in check_structure (issue #132)#138
Nas01010101 wants to merge 1 commit into
google-deepmind:mainfrom
Nas01010101:fix/issue-132-nonstr-axis-names

Conversation

@Nas01010101
Copy link
Copy Markdown

AxisName is defined as Hashable (named_axes.py), so any hashable is a valid axis name. TmpPosAxisMarker instances are commonly used as non-string batch axis names.

Two assertions in shapecheck.py incorrectly restricted axis keys to strings only:

  • _named_inline_multidimvars(): assert isinstance(subkey, str)
  • final solution verification: assert isinstance(name[1], str)

Both assertions precede code that only performs dict operations (hashable-compatible), so removing them is safe. Also update the DimVar.name and KnownDim.name type annotations to use AxisName instead of str | int.

Adds regression tests covering:

  • check_structure with known_vars containing non-string MultiDimVar keys
  • Linear layer called with a TmpPosAxisMarker batch axis (exact repro)

…mind#132)

AxisName is defined as Hashable (named_axes.py), so any hashable is a
valid axis name. TmpPosAxisMarker instances are commonly used as
non-string batch axis names.

Two assertions in shapecheck.py incorrectly restricted axis keys to
strings only:
- _named_inline_multidimvars(): assert isinstance(subkey, str)
- final solution verification: assert isinstance(name[1], str)

Both assertions precede code that only performs dict operations
(hashable-compatible), so removing them is safe. Also update the
DimVar.name and KnownDim.name type annotations to use AxisName
instead of str | int.

Adds regression tests covering:
- check_structure with known_vars containing non-string MultiDimVar keys
- Linear layer called with a TmpPosAxisMarker batch axis (exact repro)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant