Add @dca.dataclass_array decorator to customize dca params. Change default values

PiperOrigin-RevId: 475563717

Author: Etienne Pot

CHANGELOG.md

@@ -23,6 +23,10 @@ To release a new version (e.g. from `1.0.0` -> `2.0.0`):
 
 ## [Unreleased]
 
+*   Changed: By default, dataclass_array do not cast and broadcast inputs
+    anymore.
+*   Added: `@dca.dataclass_array` to customize the `dca.DataclassArray` params
+
 ## [1.1.0] - 2022-08-15
 
 *   Added: Array types can be imported directly from `dataclass_array.typing`

dataclass_array/__init__.py

@@ -32,8 +32,8 @@
 from dataclass_array import typing
 # TODO(epot): Rename array_field -> field internally
 from dataclass_array.array_dataclass import array_field as field
+from dataclass_array.array_dataclass import dataclass_array
 from dataclass_array.array_dataclass import DataclassArray
-from dataclass_array.array_dataclass import DataclassParams
 from dataclass_array.ops import stack
 from dataclass_array.vectorization import vectorize_method
 

dataclass_array/array_dataclass.py

@@ -54,18 +54,64 @@
 class DataclassParams:
   """Params controlling the DataclassArray behavior.
 
-  Saved in `cls.__dca_params__`.
+  Set by `@dca.dataclass_array`. Saved in `cls.__dca_params__`.
 
   Attributes:
-    broadcast: If `False`, disable input broadcasting
-    cast_dtype: If `False`, do not auto-cast inputs `dtype`
-    cast_list: If `False`, do not auto-cast lists to `xnp.ndarray`
+    broadcast: If `True`, enable input broadcasting
+    cast_dtype: If `True`, auto-cast inputs `dtype`
+    cast_list: If `True`, auto-cast lists to `xnp.ndarray`
   """
-  broadcast: bool = True
-  cast_dtype: bool = True
+  # If modifying this, make sure to modify `@dataclass_array` too!
+  broadcast: bool = False
+  cast_dtype: bool = False
   cast_list: bool = True
 
 
+def dataclass_array(
+    *,
+    # If modifying this, make sure to modify `DataclassParams` too!
+    broadcast: bool = False,
+    cast_dtype: bool = False,
+    cast_list: bool = True,
+) -> Callable[[type[_DcT]], type[_DcT]]:
+  """Optional decorator to customize `dca.DataclassArray` params.
+
+  Usage:
+
+  ```python
+  @dca.dataclass_array()
+  @dataclasses.dataclass(frozen=True)
+  class MyDataclass(dca.DataclassArray):
+    ...
+  ```
+
+  This decorator has to be added in addition of inheriting from
+  `dca.DataclassArray`.
+
+  Args:
+    broadcast: If `True`, enable input broadcasting
+    cast_dtype: If `True`, auto-cast inputs `dtype`
+    cast_list: If `True`, auto-cast lists to `xnp.ndarray`
+
+  Returns:
+    decorator: The decorator which will apply the options to the dataclass
+  """
+
+  def decorator(cls):
+    if not issubclass(cls, DataclassArray):
+      raise TypeError(
+          '`@dca.dataclass_array` can only be applied on `dca.DataclassArray`. '
+          f'Got: {cls}')
+    cls.__dca_params__ = DataclassParams(
+        broadcast=broadcast,
+        cast_dtype=cast_dtype,
+        cast_list=cast_list,
+    )
+    return cls
+
+  return decorator
+
+
 class MetaDataclassArray(type):
   """DataclassArray metaclass."""
 
@@ -545,8 +591,9 @@ def _broadcast_field(f: _ArrayField) -> None:
         return
       elif not self.__dca_params__.broadcast:  # Broadcasing disabled
         raise ValueError(
-            f'{type(self).__qualname__} has `broadcast=False`. Cannot '
-            f'broadcast {f.name} from {f.full_shape} to {final_shape}')
+            f'{type(self).__qualname__} has `broadcast=False`. '
+            f'Cannot broadcast {f.name} from {f.full_shape} to {final_shape}. '
+            f'To enable broadcast, use `@dca.dataclass_array(broadcast=True)`.')
       self._setattr(f.name, f.broadcast_to(final_shape))
 
     self._map_field(

dataclass_array/array_dataclass_test.py

@@ -34,24 +34,28 @@
 # TODO(epot): Test dtype `complex`, `str`
 
 
+@dca.dataclass_array(broadcast=True, cast_dtype=True)
 @dataclasses.dataclass(frozen=True)
 class Point(dca.DataclassArray):
   x: f32['*shape']
   y: f32['*shape']
 
 
+@dca.dataclass_array(broadcast=True, cast_dtype=True)
 @dataclasses.dataclass(frozen=True)
 class PointWrapper(dca.DataclassArray):
   pts: Point
   rgb: f32['*shape 3']
 
 
+@dca.dataclass_array(broadcast=True, cast_dtype=True)
 @dataclasses.dataclass(frozen=True)
 class Isometrie(dca.DataclassArray):
   r: f32['... 3 3']
   t: i32[..., 2]
 
 
+@dca.dataclass_array(broadcast=True, cast_dtype=True)
 @dataclasses.dataclass(frozen=True)
 class Nested(dca.DataclassArray):
   # pytype: disable=annotation-type-mismatch
@@ -61,6 +65,7 @@ class Nested(dca.DataclassArray):
   # pytype: enable=annotation-type-mismatch
 
 
+@dca.dataclass_array(broadcast=True, cast_dtype=True)
 @dataclasses.dataclass(frozen=True)
 class WithStatic(dca.DataclassArray):
   """Mix of static and array fields."""
@@ -541,8 +546,6 @@ def test_dataclass_params_no_cast(xnp: enp.NpModule):
 
   @dataclasses.dataclass(frozen=True)
   class PointNoCast(dca.DataclassArray):
-    __dca_params__ = dca.DataclassParams(cast_dtype=False)
-
     x: FloatArray['*shape']
     y: IntArray['*shape']
 
@@ -564,10 +567,9 @@ class PointNoCast(dca.DataclassArray):
 @enp.testing.parametrize_xnp()
 def test_dataclass_params_no_list(xnp: enp.NpModule):
 
+  @dca.dataclass_array(cast_list=False)
   @dataclasses.dataclass(frozen=True)
   class PointNoList(dca.DataclassArray):
-    __dca_params__ = dca.DataclassParams(cast_list=False)
-
     x: FloatArray['*shape']
     y: IntArray['*shape']
 
@@ -583,15 +585,13 @@ def test_dataclass_params_no_broadcast(xnp: enp.NpModule):
 
   @dataclasses.dataclass(frozen=True)
   class PointNoBroadcast(dca.DataclassArray):
-    __dca_params__ = dca.DataclassParams(broadcast=False)
-
     x: FloatArray['*shape']
     y: IntArray['*shape']
 
   with pytest.raises(ValueError, match='Cannot broadcast'):
     PointNoBroadcast(
         x=xnp.array(1, dtype=np.float16),
-        y=xnp.array([1, 2, 3], dtype=np.float16),
+        y=xnp.array([1, 2, 3], dtype=np.int32),
     )