Accept named dim in DataclassArray FloatArray['... h w 3'] (without consistency checking)

PiperOrigin-RevId: 476835750

Author: Conchylicultor

CHANGELOG.md

@@ -25,6 +25,9 @@ To release a new version (e.g. from `1.0.0` -> `2.0.0`):
 
 *   Changed: By default, dataclass_array do not cast and broadcast inputs
     anymore.
+*   Changed: `dca.DataclassArray` fields can be annotated with named axis (e.g.
+    `FloatArray['*shape h w 3']`). Note that consistency across fields is not
+    checked yet.
 *   Added: `@dca.dataclass_array` to customize the `dca.DataclassArray` params
 
 ## [1.1.0] - 2022-08-15

dataclass_array/shape_parsing.py

@@ -88,13 +88,16 @@ def get_inner_shape(shape_str: str) -> _Shape:
   shape = parser.parse(shape_str)
 
   # TODO(epot): Reraise typing with `shape` debug message
-  # TODO(epot): Support `_` & `None` dim
   if not shape or not isinstance(shape[0], _VarDim):
     raise ValueError(
         f'Shape {shape_str!r} should start by `...` or `*shape` (e.g. '
         '`f32[\'*shape 3\']`)')
 
   inner_shape = shape[1:]
+  # Currently, `_NamedDim` are accepted but consistency isn't checked across
+  # fields
+  inner_shape = tuple(
+      None if isinstance(s, _NamedDim) else s for s in inner_shape)
   if not all(isinstance(dim, (int, type(None))) for dim in inner_shape):
     raise ValueError('Only static or None dimensions supported.')
 

dataclass_array/shape_parsing_test.py

@@ -67,6 +67,7 @@ def test_parse_shape_types():
         ('... 3 5', (3, 5)),
         ('... 3 5 7', (3, 5, 7)),
         ('... 3 _ _ 7', (3, None, None, 7)),
+        ('... 3 h w 7', (3, None, None, 7)),
         ('*shape', ()),
         ('*shape _', (None,)),
         ('*shape 3', (3,)),
@@ -92,11 +93,7 @@ def test_get_inner_shape_failure_first_dim(shape_str: str):
 
 @pytest.mark.parametrize(
     'shape_str',
-    [
-        '... ...',
-        '... 3 d 1',
-        '*shape 3 d 1',
-    ],
+    ['... ...'],
 )
 def test_get_inner_shape_failure_dynamic(shape_str: str):
   with pytest.raises(ValueError, match='Only static or None dimension'):