raise warning if arrays need to be copied

Author: stes

cebra/integrations/sklearn/cebra.py

@@ -1253,6 +1253,8 @@ def transform(self,
 
         X = sklearn_utils.check_input_array(X, min_samples=len(self.offset_))
 
+        X = cebra_sklearn_dataset._ensure_writable(X)
+
         if isinstance(X, np.ndarray):
             X = torch.from_numpy(X)
 

cebra/integrations/sklearn/dataset.py

@@ -21,6 +21,8 @@
 #
 """Datasets to be used as part of the sklearn framework."""
 
+import traceback
+import warnings
 from typing import Iterable, Optional
 
 import numpy as np
@@ -34,6 +36,28 @@
 import cebra.solver
 
 
+def _ensure_writable(array: npt.NDArray) -> npt.NDArray:
+    if not array.flags.writeable:
+        stack = traceback.extract_stack()[-5:-1]
+        stack_str = ''.join(traceback.format_list(stack[-4:]))
+
+        warnings.warn(
+            ("You passed a non-writable Numpy array to CEBRA. Pytorch does currently "
+             "not support non-writable tensors. As a result, CEBRA needs to copy the "
+             "contents of the array, which might yield unnecessary memory overhead. "
+             "Ideally, adapt the code such that the array you pass to CEBRA is writable "
+             "to make your code memory efficient. "
+             "You can find more context and the rationale for this fix here: "
+             "https://github.com/AdaptiveMotorControlLab/CEBRA/pull/289."
+             "\n\n"
+             "Trace:\n" + stack_str),
+            UserWarning,
+            stacklevel=2,
+        )
+        array = array.copy()
+    return array
+
+
 class SklearnDataset(cebra.data.SingleSessionDataset):
     """Dataset for wrapping array-like input/index pairs.
 
@@ -110,9 +134,7 @@ def _parse_data(self, X: npt.NDArray):
         # one sample is a conservative default here to ensure that sklearn tests
         # passes with the correct error messages.
         X = cebra_sklearn_utils.check_input_array(X, min_samples=2)
-        # Ensure array is writable (pandas 3.0+ may return read-only arrays)
-        if not X.flags.writeable:
-            X = X.copy()
+        X = _ensure_writable(X)
         self.neural = torch.from_numpy(X).float().to(self.device)
 
     def _parse_labels(self, labels: Optional[tuple]):
@@ -146,11 +168,10 @@ def _parse_labels(self, labels: Optional[tuple]):
                     f"or lists that can be converted to arrays, but got {type(y)}"
                 )
 
+            y = _ensure_writable(y)
+
             # Define the index as either continuous or discrete indices, depending
             # on the dtype in the index array.
-            # Ensure array is writable (pandas 3.0+ may return read-only arrays)
-            if not y.flags.writeable:
-                y = y.copy()
             if cebra.helper._is_floating(y):
                 y = torch.from_numpy(y).float()
                 if y.dim() == 1:

tests/test_sklearn.py

@@ -1544,3 +1544,29 @@ def test_last_incomplete_batch_smaller_than_offset():
     model.fit(train.neural, train.continuous)
 
     _ = model.transform(train.neural, batch_size=300)
+
+
+def test_non_writable_array():
+    import numpy as np
+    import pytest
+
+    import cebra
+
+    # Create a numpy array and make it non-writable
+    X = np.random.randn(100, 10)
+    y = np.random.randn(100, 2)
+
+    X.setflags(write=False)
+    y.setflags(write=False)
+
+    with pytest.raises(ValueError, match="assignment destination is read-only"):
+        X[:] = 0
+        y[:] = 0
+
+    cebra_model = cebra.CEBRA(max_iterations=2, batch_size=32, device="cpu")
+
+    # This should not raise an exception even though arrays are not writable
+    cebra_model.fit(X, y)
+    embedding = cebra_model.transform(X)
+    assert isinstance(embedding, np.ndarray)
+    assert embedding.shape[0] == X.shape[0]