Add dca.concatenate to dataclass_array

Allows dataclass arrays to be not only stacked, but also concatenated together.

PiperOrigin-RevId: 546279663

Author: PeterZhizhin

CHANGELOG.md

@@ -23,6 +23,10 @@ To release a new version (e.g. from `1.0.0` -> `2.0.0`):
 
 ## [Unreleased]
 
+## [1.4.3] - 2023-07-07
+
+* Add `dca.concat` method in addition to `dca.stack`.
+
 ## [1.4.2] - 2023-06-29
 
 *   Now require Python 3.9 (drop 3.8 support)
@@ -69,8 +73,9 @@ To release a new version (e.g. from `1.0.0` -> `2.0.0`):
 
 *   Initial release
 
-[Unreleased]: https://github.com/google-research/dataclass_array/compare/v1.4.2...HEAD
-[1.4.1]: https://github.com/google-research/dataclass_array/compare/v1.4.1...v1.4.2
+[Unreleased]: https://github.com/google-research/dataclass_array/compare/v1.4.3...HEAD
+[1.4.3]: https://github.com/google-research/dataclass_array/compare/v1.4.2...v1.4.3
+[1.4.2]: https://github.com/google-research/dataclass_array/compare/v1.4.1...v1.4.2
 [1.4.1]: https://github.com/google-research/dataclass_array/compare/v1.4.0...v1.4.1
 [1.4.0]: https://github.com/google-research/dataclass_array/compare/v1.3.0...v1.4.0
 [1.3.0]: https://github.com/google-research/dataclass_array/compare/v1.2.1...v1.3.0

dataclass_array/__init__.py

@@ -34,6 +34,7 @@
 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.ops import concat
 from dataclass_array.ops import stack
 from dataclass_array.vectorization import vectorize_method
 

dataclass_array/array_dataclass_test.py

@@ -734,6 +734,23 @@ class PointDynamicShape(dca.DataclassArray):
     )
 
 
+@enp.testing.parametrize_xnp()
+@pytest.mark.parametrize('batch_shape', [(1,), (3,)])
+def test_concatenate(xnp: enp.NpModule, batch_shape: Shape):
+  class TestConcatenateClass(dca.DataclassArray):
+    x: FloatArray['*shape 3']
+    y: FloatArray['*shape']
+
+  p = TestConcatenateClass(
+      x=xnp.zeros(batch_shape + (3,), dtype=xnp.float32),
+      y=xnp.zeros(batch_shape, dtype=xnp.float32),
+  )
+
+  p_concatenated = dca.concat([p, p, p])
+  assert p_concatenated.x.shape == tuple(x * 3 for x in batch_shape) + (3,)
+  assert p_concatenated.y.shape == tuple(x * 3 for x in batch_shape)
+
+
 def test_class_getitem():
   assert Point == Point  # pylint: disable=comparison-with-itself
   assert Point[''] == Point['']

dataclass_array/ops.py

@@ -16,20 +16,40 @@
 
 from __future__ import annotations
 
-from typing import Iterable  # pylint: disable=g-multiple-import
+import functools
+from typing import Any, Callable, Iterable, Optional  # pylint: disable=g-multiple-import
 
 from dataclass_array import array_dataclass
-from dataclass_array.typing import DcT
+from dataclass_array.typing import Array, DcT  # pylint: disable=g-multiple-import
 from dataclass_array.utils import np_utils
+from etils import enp
 from etils import epy
 
 
-def stack(
-    arrays: Iterable[DcT],  # list[_DcT['*shape']]
+def _ops_base(
+    arrays: Iterable[DcT],
     *,
-    axis: int = 0,
-) -> DcT:  # _DcT['len(arrays) *shape']:
-  """Stack dataclasses together."""
+    axis: int,
+    array_fn: Callable[
+        [
+            enp.NpModule,
+            int,
+            Any,  # array_dataclass._ArrayField[Array['*din']],
+        ],
+        Array['*dout'],
+    ],
+    dc_fn: Optional[
+        Callable[
+            [
+                enp.NpModule,
+                int,
+                Any,  # array_dataclass._ArrayField[DcT],
+            ],
+            DcT,
+        ]
+    ],
+) -> DcT:
+  """Base function for all ops."""
   arrays = list(arrays)
   first_arr = arrays[0]
 
@@ -61,9 +81,41 @@ def stack(
   # jax.tree_map(lambda x, y: x+y, (None, 10), (1, 2)) == (None, 12)
   # Similarly, static values will be the ones from the first element.
   merged_arr = first_arr._map_field(  # pylint: disable=protected-access
-      array_fn=lambda f: xnp.stack(  # pylint: disable=g-long-lambda
+      array_fn=functools.partial(array_fn, xnp, axis),
+      dc_fn=functools.partial(dc_fn, xnp, axis),
+  )
+  return merged_arr
+
+
+def stack(
+    arrays: Iterable[DcT],  # list[_DcT['*shape']]
+    *,
+    axis: int = 0,
+) -> DcT:  # _DcT['len(arrays) *shape']:
+  """Stack dataclasses together."""
+  return _ops_base(
+      arrays,
+      axis=axis,
+      array_fn=lambda xnp, axis, f: xnp.stack(  # pylint: disable=g-long-lambda
           [getattr(arr, f.name) for arr in arrays], axis=axis
       ),
-      dc_fn=lambda f: stack([getattr(arr, f.name) for arr in arrays]),
+      dc_fn=lambda xnp, axis, f: stack(  # pylint: disable=g-long-lambda
+          [getattr(arr, f.name) for arr in arrays],
+          axis=axis,
+      ),
+  )
+
+
+def concat(arrays: Iterable[DcT], *, axis: int = 0) -> DcT:
+  """Concatenate dataclasses together."""
+  return _ops_base(
+      arrays,
+      axis=axis,
+      array_fn=lambda xnp, axis, f: xnp.concatenate(  # pylint: disable=g-long-lambda
+          [getattr(arr, f.name) for arr in arrays], axis=axis
+      ),
+      dc_fn=lambda xnp, axis, f: concat(  # pylint: disable=g-long-lambda
+          [getattr(arr, f.name) for arr in arrays],
+          axis=axis,
+      ),
   )
-  return merged_arr