Support typing annotation shape for dataclass array (e.g. ray: Ray['*batch_shape _ _'])

PiperOrigin-RevId: 466946356

Author: Conchylicultor

CHANGELOG.md

@@ -23,11 +23,13 @@ To release a new version (e.g. from `1.0.0` -> `2.0.0`):
 
 ## [Unreleased]
 
+*   Added: Syntax to specify the shape of the DataclassArray (e.g. `MyRay['h
+    w']`).
+
 ## [1.0.0] - 2022-08-08
 
 *   Initial release
 
-
 [Unreleased]: https://github.com/google-research/dataclass_array/compare/v1.0.0...HEAD
 [1.0.0]: https://github.com/google-research/dataclass_array/compare/v0.1.0...v1.0.0
 [0.1.0]: https://github.com/google-research/dataclass_array/releases/tag/v0.1.0

README.md

@@ -77,15 +77,18 @@ A `DataclassArray` has 2 types of fields:
 class MyArray(dca.DataclassArray):
   # Array fields
   a: FloatArray['*batch_shape 3']  # Defined by `etils.array_types`
-  b: Ray  # Nested DataclassArray (inner shape == `()`)
+  b: FloatArray['*batch_shape _ _']  # Dynamic shape
+  c: Ray  # Nested DataclassArray (equivalent to `Ray['*batch_shape']`)
+  d: Ray['*batch_shape 6']
 
   # Array fields explicitly defined
-  c: Any = dca.field(shape=(3,), dtype=np.float32)
-  d: Ray = dca.field(shape=(3,), dtype=Ray)  # Nested DataclassArray
+  e: Any = dca.field(shape=(3,), dtype=np.float32)
+  f: Any = dca.field(shape=(None,  None), dtype=np.float32)  # Dynamic shape
+  g: Ray = dca.field(shape=(3,), dtype=Ray)  # Nested DataclassArray
 
   # Static field (everything not defined as above)
-  e: float
-  f: np.array
+  static0: float
+  static1: np.array
 ```
 
 ### Vectorization

dataclass_array/array_dataclass.py

@@ -20,9 +20,10 @@
 import typing
 from typing import Any, Callable, ClassVar, Generic, Iterator, Optional, Tuple, Type, TypeVar, Union
 
+from dataclass_array import field_utils
 from dataclass_array import shape_parsing
 from dataclass_array import type_parsing
-from dataclass_array.typing import Axes, DcOrArray, DcOrArrayT, DTypeArg, Shape  # pylint: disable=g-multiple-import
+from dataclass_array.typing import Axes, DcOrArray, DcOrArrayT, DTypeArg, DynamicShape, Shape  # pylint: disable=g-multiple-import
 from dataclass_array.utils import np_utils
 from dataclass_array.utils import py_utils
 import einops
@@ -32,7 +33,8 @@
 from etils import epy
 from etils.array_types import Array
 import numpy as np
-from typing_extensions import Literal, TypeAlias  # pylint: disable=g-multiple-import
+import typing_extensions
+from typing_extensions import Annotated, Literal, TypeAlias  # pylint: disable=g-multiple-import
 
 
 lazy = enp.lazy
@@ -65,7 +67,16 @@ class DataclassParams:
   cast_list: bool = True
 
 
-class DataclassArray:
+class MetaDataclassArray(type):
+  """DataclassArray metaclass."""
+
+  # TODO(b/204422756): We cannot use `__class_getitem__` due to b/204422756
+  def __getitem__(cls, spec):
+    # Not clear how this would interact if cls is also a `Generic`
+    return Annotated[cls, field_utils.ShapeAnnotation(spec)]
+
+
+class DataclassArray(metaclass=MetaDataclassArray):
   """Dataclass which behaves like an array.
 
   Usage:
@@ -385,7 +396,7 @@ def _all_array_fields(self) -> dict[str, _ArrayField]:
     if cls._dca_fields_metadata is None:  # pylint: disable=protected-access
       # At this point, `ForwardRef` should have been resolved.
       try:
-        hints = typing.get_type_hints(cls)
+        hints = typing_extensions.get_type_hints(cls, include_extras=True)
       except Exception as e:  # pylint: disable=broad-except
         epy.reraise(
             e,
@@ -706,7 +717,7 @@ class _ArrayFieldMetadata:
     dtype: Type of the array. Can be `array_types.dtypes.DType` or
       `dca.DataclassArray` for nested arrays.
   """
-  inner_shape_non_static: Shape
+  inner_shape_non_static: DynamicShape
   dtype: Union[array_types.dtypes.DType, type[DataclassArray]]
 
   def __post_init__(self):
@@ -739,7 +750,7 @@ class _ArrayField(_ArrayFieldMetadata, Generic[DcOrArrayT]):
     host: Dataclass instance who this field is attached too
   """
   name: str
-  host: DataclassArray
+  host: DataclassArray = dataclasses.field(repr=False)
 
   @property
   def qualname(self) -> str:
@@ -851,18 +862,18 @@ def _make_field_metadata(
 def _type_to_field_metadata(hint: TypeAlias) -> Optional[_ArrayFieldMetadata]:
   """Converts type hint to extract `inner_shape`, `dtype`."""
   array_type = type_parsing.get_array_type(hint)
-  if isinstance(array_type, type) and issubclass(array_type, DataclassArray):
-    # TODO(epot): Should support `ray: Ray[..., 3]` ?
-    return _ArrayFieldMetadata(inner_shape_non_static=(), dtype=array_type)
+
+  if isinstance(array_type, field_utils.DataclassWithShape):
+    dtype = array_type.cls
   elif isinstance(array_type, array_types.ArrayAliasMeta):
-    assert array_type is not None
-    try:
-      return _ArrayFieldMetadata(
-          inner_shape_non_static=shape_parsing.get_inner_shape(
-              array_type.shape),
-          dtype=array_type.dtype,
-      )
-    except Exception as e:  # pylint: disable=broad-except
-      epy.reraise(e, prefix=f'Invalid shape annotation {hint}.')
+    dtype = array_type.dtype
   else:  # Not a supported type: Static field
     return None
+
+  try:
+    return _ArrayFieldMetadata(
+        inner_shape_non_static=shape_parsing.get_inner_shape(array_type.shape),
+        dtype=dtype,
+    )
+  except Exception as e:  # pylint: disable=broad-except
+    epy.reraise(e, prefix=f'Invalid shape annotation {hint}.')

dataclass_array/array_dataclass_test.py

@@ -56,7 +56,7 @@ class Isometrie(dca.DataclassArray):
 class Nested(dca.DataclassArray):
   # pytype: disable=annotation-type-mismatch
   iso: Isometrie
-  iso_batched: Isometrie = dca.field(shape=(3, 7), dtype=Isometrie)
+  iso_batched: Isometrie['*batch_shape 3 7']
   pt: Point = dca.field(shape=(3,), dtype=Point)
   # pytype: enable=annotation-type-mismatch
 
@@ -645,3 +645,12 @@ class PointDynamicShape(dca.DataclassArray):
         x=xnp.zeros(batch_shape + (2, 3), dtype=np.float32),
         y=xnp.zeros(batch_shape + (2, 1), dtype=np.int32),  # < 2 != 3
     )
+
+
+def test_class_getitem():
+  assert Point == Point  # pylint: disable=comparison-with-itself
+  assert Point[''] == Point['']
+  assert Point['h w'] == Point['h w']
+  assert Point[''] != Point
+  assert Point[''] != Point['h w']
+  assert Point[''] != Isometrie['']

dataclass_array/field_utils.py

@@ -0,0 +1,67 @@
+# Copyright 2022 The dataclass_array 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.
+
+"""Field utils."""
+
+from __future__ import annotations
+
+import dataclasses
+
+from dataclass_array import array_dataclass
+from dataclass_array.typing import TypeAlias
+from etils import epy
+from etils.array_types import Array
+import typing_extensions
+
+
+@dataclasses.dataclass(eq=True, frozen=True)
+class ShapeAnnotation:
+  """Annotations for `Ray[''] == Annotated[Ray, ShapeAnnotation('')]`."""
+  shape: str
+
+  def __post_init__(self):
+    # Normalize shape
+    # Might be a cleaner way to do this
+    super().__setattr__('shape', Array[self.shape].shape)
+
+
+@dataclasses.dataclass(eq=True, frozen=True)
+class DataclassWithShape:
+  """Structure which represent `Ray['h w']`."""
+  cls: type[array_dataclass.DataclassArray]
+  shape: str
+
+  @classmethod
+  def from_hint(cls, hint: TypeAlias) -> DataclassWithShape:
+    """Factory to create the `DataclassWithShape` from `MyDca['h w']`."""
+    assert cls.is_dca(hint)
+
+    # Extract the shape
+    shape = '...'
+    if typing_extensions.get_origin(hint) is typing_extensions.Annotated:
+      shapes = [a for a in hint.__metadata__ if isinstance(a, ShapeAnnotation)]  # pytype: disable=attribute-error
+      if len(shapes) > 1:
+        raise ValueError(f'Conflicting annotations for {hint}')
+      elif len(shapes) == 1:
+        (shape,) = shapes
+        shape = shape.shape
+
+      hint = hint.__origin__
+    return cls(cls=hint, shape=shape)
+
+  @classmethod
+  def is_dca(cls, hint: TypeAlias) -> bool:
+    if typing_extensions.get_origin(hint) is typing_extensions.Annotated:
+      hint = hint.__origin__
+    return epy.issubclass(hint, array_dataclass.DataclassArray)

dataclass_array/shape_parsing.py

@@ -91,7 +91,8 @@ def get_inner_shape(shape_str: str) -> _Shape:
   # TODO(epot): Support `_` & `None` dim
   if not shape or not isinstance(shape[0], _VarDim):
     raise ValueError(
-        'Shape should start by `...` or `*shape` (e.g. `f32[\'*shape 3\']`)')
+        f'Shape {shape_str!r} should start by `...` or `*shape` (e.g. '
+        '`f32[\'*shape 3\']`)')
 
   inner_shape = shape[1:]
   if not all(isinstance(dim, (int, type(None))) for dim in inner_shape):

dataclass_array/shape_parsing_test.py

@@ -86,7 +86,7 @@ def test_get_inner_shape(shape_str, shape_tuple):
     ],
 )
 def test_get_inner_shape_failure_first_dim(shape_str: str):
-  with pytest.raises(ValueError, match='Shape should start'):
+  with pytest.raises(ValueError, match='Shape .* should start'):
     shape_parsing.get_inner_shape(shape_str)
 
 

dataclass_array/type_parsing.py

@@ -22,10 +22,11 @@
 from typing import Any, Callable, Optional
 
 from dataclass_array import array_dataclass
+from dataclass_array import field_utils
+from dataclass_array.typing import TypeAlias
 from etils import array_types as array_types_lib
 import typing_extensions  # TODO(py38): Remove
 
-TypeAlias = Any  # TODO(py39): Use real alias once 3.7 is dropped.
 _LeafFn = Callable[[TypeAlias], None]
 
 _NoneType = type(None)
@@ -74,15 +75,15 @@ def _collect_leaf_types(hint):
   return all_types
 
 
-def get_array_type(hint: TypeAlias) -> Optional[type[Any]]:
+def get_array_type(hint: TypeAlias) -> Optional[Any]:
   """Returns the array type, or `None` if no type was detected.
 
   Example:
 
   ```python
   get_array_type(f32[..., 3]) -> f32[..., 3]
-  get_array_type(dca.Ray) -> dca.Ray
-  get_array_type(Optional[dca.Ray]) -> dca.Ray
+  get_array_type(dca.Ray) -> dca.Ray['...']
+  get_array_type(Optional[dca.Ray]) -> dca.Ray['...']
   get_array_type(dca.Ray | dca.Camera | None) -> dca.DataclassArray
   get_array_type(Any) -> None  # Any not an array type
   get_array_type(dca.Ray | int) -> None  # int not an array type
@@ -106,9 +107,8 @@ def get_array_type(hint: TypeAlias) -> Optional[type[Any]]:
   array_types = []
   other_types = []
   for leaf in leaf_types:
-    if (isinstance(leaf, type) and
-        issubclass(leaf, array_dataclass.DataclassArray)):
-      dc_types.append(leaf)
+    if field_utils.DataclassWithShape.is_dca(leaf):
+      dc_types.append(field_utils.DataclassWithShape.from_hint(leaf))
     elif isinstance(leaf, array_types_lib.ArrayAliasMeta):
       array_types.append(leaf)
     else:
@@ -122,7 +122,17 @@ def get_array_type(hint: TypeAlias) -> Optional[type[Any]]:
         'this feature.')
   if dc_types:
     if len(dc_types) > 1:
-      return array_dataclass.DataclassArray
+      # Validate the inner shape
+      common_shapes = {x.shape for x in dc_types}
+      if len(common_shapes) != 1:
+        raise NotImplementedError(
+            f'{hint} mix dataclass with different inner shape. Please open an '
+            'issue if you need this feature.')
+      (common_shape,) = common_shapes
+      return field_utils.DataclassWithShape(
+          cls=array_dataclass.DataclassArray,
+          shape=common_shape,
+      )
     else:
       return dc_types[0]
   if array_types:

dataclass_array/type_parsing_test.py

@@ -25,6 +25,7 @@
 import pytest
 
 
+_DS = dca.field_utils.DataclassWithShape
 Ray = dca.testing.Ray
 
 
@@ -39,7 +40,10 @@ class Camera(dca.DataclassArray):
     [
         (int, [int]),
         (Ray, [Ray]),
+        (Ray['h w'], [Ray['h w']]),
+        (Ray[..., 3], [Ray[..., 3]]),
         (Union[Ray, int], [Ray, int]),
+        (Union[Ray['h w'], int], [Ray['h w'], int]),
         (Union[Ray, int, None], [Ray, int, None]),
         (Optional[Ray], [Ray, None]),
         (Optional[Union[Ray, int]], [Ray, int, None]),
@@ -55,10 +59,12 @@ def test_get_leaf_types(hint, expected):
     'hint, expected',
     [
         (int, None),
-        (Ray, Ray),
-        (Optional[Ray], Ray),
-        (Union[Ray, Camera], dca.DataclassArray),
-        (Union[Ray, Camera, None], dca.DataclassArray),
+        (Ray, _DS(Ray, '...')),
+        (Ray['h w'], _DS(Ray, 'h w')),
+        (Ray[..., 3], _DS(Ray, '... 3')),
+        (Optional[Ray], _DS(Ray, '...')),
+        (Union[Ray, Camera], _DS(dca.DataclassArray, '...')),
+        (Union[Ray, Camera, None], _DS(dca.DataclassArray, '...')),
         (Union[Ray, int], None),
         (Union[Ray, int, None], None),
         (Union[f32[3, 3], int, None], None),
@@ -72,9 +78,51 @@ def test_get_array_type(hint, expected):
   assert type_parsing.get_array_type(hint) == expected
 
 
+@pytest.mark.parametrize(
+    'hint, expected',
+    [
+        (Ray, _DS(Ray, '...')),
+        (Ray['h w'], _DS(Ray, 'h w')),
+        (Ray[..., 3], _DS(Ray, '... 3')),
+    ],
+)
+def test_from_hint(hint, expected):
+  assert dca.field_utils.DataclassWithShape.from_hint(hint) == expected
+
+
 def test_get_array_type_error():
   with pytest.raises(NotImplementedError):
     type_parsing.get_array_type(Union[Ray, f32[3, 3]])
 
   with pytest.raises(NotImplementedError):
     type_parsing.get_array_type(Union[FloatArray[..., 3], f32[3, 3]])
+
+
+@pytest.mark.parametrize(
+    'hint, expected',
+    [
+        (
+            Ray,
+            dca.array_dataclass._ArrayFieldMetadata(
+                inner_shape_non_static=(),
+                dtype=Ray,
+            ),
+        ),
+        (
+            Ray[..., 3],
+            dca.array_dataclass._ArrayFieldMetadata(
+                inner_shape_non_static=(3,),
+                dtype=Ray,
+            ),
+        ),
+        (
+            Ray['*shape 4 _'],
+            dca.array_dataclass._ArrayFieldMetadata(
+                inner_shape_non_static=(4, None),
+                dtype=Ray,
+            ),
+        ),
+    ],
+)
+def test_type_to_field_metadata(hint, expected):
+  assert dca.array_dataclass._type_to_field_metadata(hint) == expected