Fix read-only observation arrays in transforms

ax/adapter/transforms/bilog_y.py

@@ -102,6 +102,10 @@ def _reusable_transform(
         transform: Callable[[npt.NDArray, npt.NDArray], npt.NDArray],
     ) -> list[ObservationData]:
         for obsd in observation_data:
+            if not obsd.means.flags.writeable:
+                obsd.means = obsd.means.copy()
+            if not obsd.covariance.flags.writeable:
+                obsd.covariance = obsd.covariance.copy()
             for i, m in enumerate(obsd.metric_signatures):
                 if m in self.metric_to_bound.keys():
                     bound = self.metric_to_bound[m]

ax/adapter/transforms/log_y.py

@@ -174,6 +174,10 @@ def _tf_obs_data(
         if len(self.metric_signatures) == 0:
             return observation_data
         for obsd in observation_data:
+            if not obsd.means.flags.writeable:
+                obsd.means = obsd.means.copy()
+            if not obsd.covariance.flags.writeable:
+                obsd.covariance = obsd.covariance.copy()
             cov = obsd.covariance
             # Check for correlations in the covariance matrix
             diff = cov - np.diag(np.diag(cov))

ax/adapter/transforms/power_transform_y.py

@@ -111,6 +111,10 @@ def _transform_observation_data(
     ) -> list[ObservationData]:
         """Winsorize observation data in place."""
         for obsd in observation_data:
+            if not obsd.means.flags.writeable:
+                obsd.means = obsd.means.copy()
+            if not obsd.covariance.flags.writeable:
+                obsd.covariance = obsd.covariance.copy()
             for i, m in enumerate(obsd.metric_signatures):
                 if m in self.metric_signatures:
                     transform = self.power_transforms[m].transform
@@ -130,6 +134,10 @@ def _untransform_observation_data(
     ) -> list[ObservationData]:
         """Winsorize observation data in place."""
         for obsd in observation_data:
+            if not obsd.means.flags.writeable:
+                obsd.means = obsd.means.copy()
+            if not obsd.covariance.flags.writeable:
+                obsd.covariance = obsd.covariance.copy()
             for i, m in enumerate(obsd.metric_signatures):
                 if m in self.metric_signatures:
                     l, u = self.inv_bounds[m]

ax/adapter/transforms/standardize_y.py

@@ -79,6 +79,8 @@ def _transform_observation_data(
             means = np.array([self.Ymean[m] for m in obsd.metric_signatures])
             stds = np.array([self.Ystd[m] for m in obsd.metric_signatures])
             obsd.means = (obsd.means - means) / stds
+            if not obsd.covariance.flags.writeable:
+                obsd.covariance = obsd.covariance.copy()
             obsd.covariance /= np.dot(stds[:, None], stds[:, None].transpose())
         return observation_data
 
@@ -221,6 +223,8 @@ def _untransform_observation_data(
             means = np.array([self.Ymean[m] for m in obsd.metric_signatures])
             stds = np.array([self.Ystd[m] for m in obsd.metric_signatures])
             obsd.means = obsd.means * stds + means
+            if not obsd.covariance.flags.writeable:
+                obsd.covariance = obsd.covariance.copy()
             obsd.covariance *= np.dot(stds[:, None], stds[:, None].transpose())
         return observation_data
 

ax/adapter/transforms/tests/test_winsorize_transform.py

@@ -25,6 +25,7 @@
 from ax.core.arm import Arm
 from ax.core.data import Data
 from ax.core.metric import Metric
+from ax.core.observation import Observation, ObservationData
 from ax.core.objective import MultiObjective, Objective, ScalarizedObjective
 from ax.core.optimization_config import (
     MultiObjectiveOptimizationConfig,
@@ -180,6 +181,24 @@ def test_TransformObservations(self) -> None:
             tf_means = [obsd.means.tolist() for obsd in observation_data]
             self.assertListEqual(tf_means, expected)
 
+    def test_transform_observations_with_read_only_means(self) -> None:
+        means = np.array([0.0, 3.0])
+        covariance = np.eye(2)
+        means.flags.writeable = False
+        covariance.flags.writeable = False
+        observation = Observation(
+            features=self.observations[0].features,
+            data=ObservationData(
+                metric_signatures=["m1", "m2"],
+                means=means,
+                covariance=covariance,
+            ),
+        )
+
+        transformed = self.t.transform_observations([observation])
+
+        self.assertEqual(transformed[0].data.means.tolist(), [0.0, 2.0])
+
     def test_optimization_config_default(self) -> None:
         # Specify the winsorization
         experiment = get_experiment_with_observations(

ax/adapter/transforms/winsorize.py

@@ -140,6 +140,8 @@ def _transform_observation_data(
     ) -> list[ObservationData]:
         """Winsorize observation data in place."""
         for obsd in observation_data:
+            if not obsd.means.flags.writeable:
+                obsd.means = obsd.means.copy()
             for idx, metric_signature in enumerate(obsd.metric_signatures):
                 if metric_signature not in self.cutoffs:
                     continue