Support pickle protocol for DataclassArray.

PiperOrigin-RevId: 675147723

Author: Jan Hosang

CHANGELOG.md

@@ -23,6 +23,8 @@ To release a new version (e.g. from `1.0.0` -> `2.0.0`):
 
 ## [Unreleased]
 
+*   Add pickle support to `DataclassArray`.
+
 ## [1.5.2] - 2024-03-19
 
 *   Drop Python 3.10 support

dataclass_array/array_dataclass.py

@@ -798,6 +798,21 @@ def tree_unflatten(
       array_field_values: list[DcOrArray],
   ) -> _DcT:
     """`jax.tree_utils` support."""
+    return cls._tree_unflatten(
+        metadata=metadata,
+        array_field_values=array_field_values,
+        constructor=cls,
+    )
+
+  @classmethod
+  def _tree_unflatten(
+      cls: Type[_DcT],
+      *,
+      metadata: _TreeMetadata,
+      array_field_values: list[DcOrArray],
+      constructor: Callable[..., _DcT],
+  ) -> _DcT:
+    """Initialize a model after serialization."""
     array_field_kwargs = dict(
         zip(
             metadata.array_field_names,
@@ -813,7 +828,7 @@ def tree_unflatten(
       else:
         non_init_fields[k] = v
 
-    self = cls(**array_field_kwargs, **init_fields)
+    self = constructor(**array_field_kwargs, **init_fields)
     # Currently it's not clear how to handle non-init fields so raise an error
     if non_init_fields:
       if set(non_init_fields) - self.__dca_non_init_fields__:
@@ -844,6 +859,23 @@ def __tf_unflatten__(
   ) -> _DcT:
     return cls.tree_unflatten(metadata, components)
 
+  def __getstate__(self) -> tuple[_TreeMetadata, tuple[DcOrArray, ...]]:
+    components, metadata = self._tree_flatten()
+    return metadata, components
+
+  def __setstate__(self, state: tuple[_TreeMetadata, list[DcOrArray]]) -> None:
+
+    def constructor(**kwargs):
+      self.__init__(**kwargs)
+      return self
+
+    metadata, components = state
+    type(self)._tree_unflatten(
+        metadata=metadata,
+        array_field_values=components,
+        constructor=constructor,
+    )
+
   def _setattr(self, name: str, value: Any) -> None:
     """Like setattr, but support `frozen` dataclasses."""
     object.__setattr__(self, name, value)
@@ -1023,8 +1055,8 @@ class _ArrayFieldMetadata:
   Attributes:
     inner_shape_non_static: Inner shape. Can contain non-static dims (e.g.
       `(None, 3)`)
-    dtype: Type of the array. Can be `enp.dtypes.DType` or
-      `dca.DataclassArray` for nested arrays.
+    dtype: Type of the array. Can be `enp.dtypes.DType` or `dca.DataclassArray`
+      for nested arrays.
   """
 
   inner_shape_non_static: DynamicShape

dataclass_array/array_dataclass_test.py

@@ -16,6 +16,7 @@
 
 from __future__ import annotations
 
+import pickle
 from typing import Any
 
 import dataclass_array as dca
@@ -769,3 +770,15 @@ def test_class_getitem():
   assert Point[''] != Point
   assert Point[''] != Point['h w']
   assert Point[''] != Isometrie['']
+
+
+@enp.testing.parametrize_xnp()
+@pytest.mark.parametrize('batch_shape', [(), (1, 3)])
+def test_dataclass_pickle_unpickle(xnp: enp.NpModule, batch_shape: Shape):
+  expected = Point(
+      x=xnp.zeros(batch_shape, dtype=xnp.float32),
+      y=xnp.zeros(batch_shape, dtype=xnp.float32),
+  )
+  buffer = pickle.dumps(expected)
+  actual = pickle.loads(buffer)
+  dca.testing.assert_array_equal(actual, expected)