Add back data setter oops

PiperOrigin-RevId: 638764523

iris/normalizer.py

@@ -52,6 +52,11 @@ def merge(self, data: Dict[str, Any]) -> None:
   def data(self) -> Dict[str, Any]:
     """Returns copy of current data in buffer."""
 
+  @data.setter
+  @abc.abstractmethod
+  def data(self, new_data: Dict[str, Any]) -> None:
+    """Sets data of buffer."""
+
   @property
   @abc.abstractmethod
   def shape(self) -> Sequence[int]:
@@ -84,6 +89,10 @@ def merge(self, data: Dict[str, Any]) -> None:
   def data(self) -> Dict[str, Any]:
     return {}
 
+  @data.setter
+  def data(self, new_data: Dict[str, Any]) -> None:
+    pass
+
   @property
   def shape(self) -> Sequence[int]:
     return ()
@@ -160,6 +169,10 @@ def merge(self, data: Dict[str, Any]) -> None:
   def data(self) -> Dict[str, Any]:
     return copy.deepcopy(self._data)
 
+  @data.setter
+  def data(self, new_data: Dict[str, Any]) -> None:
+    self._data = copy.deepcopy(new_data)
+
   @property
   def shape(self) -> Sequence[int]:
     return self._shape

iris/normalizer_test.py

@@ -18,6 +18,20 @@
 from absl.testing import absltest
 
 
+class BufferTest(absltest.TestCase):
+
+  def test_meanstdbuffer(self):
+    buffer = normalizer.MeanStdBuffer((1))
+    buffer.push(np.asarray(10.0))
+    buffer.push(np.asarray(11.0))
+
+    new_buffer = normalizer.MeanStdBuffer((1))
+    new_buffer.data = buffer.data
+
+    self.assertEqual(new_buffer._std, buffer._std)
+    self.assertEqual(new_buffer._data['n'], buffer._data['n'])
+
+
 class NormalizerTest(absltest.TestCase):
 
   def test_no_normalizer(self):