From 5fd98c45bc09268cf8c737ef6e63a805fb04772b Mon Sep 17 00:00:00 2001
From: Tristan Simas <tristan.simas@mail.mcgill.ca>
Date: Wed, 7 Oct 2026 21:53:53 -0400
Subject: [PATCH] perf: normalize image planes into one owned mutable buffer

---
 src/arraybridge/__init__.py         |  2 +-
 src/arraybridge/array_operations.py | 46 +++++++++++++++++++++++++++--
 src/arraybridge/types.py            | 14 +++++++++
 tests/test_native_geometry.py       | 15 ++++++++++
 4 files changed, 74 insertions(+), 3 deletions(-)

diff --git a/src/arraybridge/__init__.py b/src/arraybridge/__init__.py
index 5985b4c..eb83370 100644
--- a/src/arraybridge/__init__.py
+++ b/src/arraybridge/__init__.py
@@ -10,7 +10,7 @@
 NumPy/numcodecs import cost.
 """
 
-__version__ = "0.3.7"
+__version__ = "0.3.8"
 
 _LAZY_EXPORTS: dict[str, str] = {
     "MemoryType": ".types",
diff --git a/src/arraybridge/array_operations.py b/src/arraybridge/array_operations.py
index b435926..90f370c 100644
--- a/src/arraybridge/array_operations.py
+++ b/src/arraybridge/array_operations.py
@@ -114,8 +114,50 @@ def broadcast_to(data: Any, shape: tuple[int, ...], module: Any) -> Any:
     def ones_like(reference: Any, shape: tuple[int, ...], dtype: Any, module: Any) -> Any:
         return module.ones(shape, dtype=dtype)
 
+    @classmethod
+    def normalize_planes(
+        cls,
+        data: Any,
+        dtype: Any,
+        scales: Sequence[float | None],
+        module: Any,
+    ) -> Any:
+        """Cast and scale each plane before assembling its owned output."""
+        return cls.stack(
+            tuple(
+                (
+                    cls.cast(plane, dtype, module)
+                    if scale is None
+                    else cls.cast(plane, dtype, module) / float(scale)
+                )
+                for plane, scale in zip(data, scales, strict=True)
+            ),
+            module,
+        )
+
+
+class MutableArrayOperations(ArrayOperations):
+    """Arrays whose allocated output supports native in-place arithmetic."""
+
+    @classmethod
+    def normalize_planes(
+        cls,
+        data: Any,
+        dtype: Any,
+        scales: Sequence[float | None],
+        module: Any,
+    ) -> Any:
+        # Integer division promotes the output dtype in the original recipe.
+        if not np.issubdtype(np.dtype(dtype), np.inexact) or len(data) == 0:
+            return super().normalize_planes(data, dtype, scales, module)
+        normalized = module.array(data, dtype=dtype, copy=True)
+        for index, scale in enumerate(scales):
+            if scale is not None:
+                normalized[index] /= float(scale)
+        return normalized
+
 
-class NumpyArrayOperations(ArrayOperations):
+class NumpyArrayOperations(MutableArrayOperations):
     """Numpy native operation leaves."""
 
     @staticmethod
@@ -148,7 +190,7 @@ def scale_dtype(result: Any, target_dtype: Any, module: Any) -> Any:
         return scaled.astype(target_dtype)
 
 
-class CupyArrayOperations(ArrayOperations):
+class CupyArrayOperations(MutableArrayOperations):
     """Cupy native operation leaves."""
 
     @staticmethod
diff --git a/src/arraybridge/types.py b/src/arraybridge/types.py
index e7c68c2..7ec4d64 100644
--- a/src/arraybridge/types.py
+++ b/src/arraybridge/types.py
@@ -804,6 +804,20 @@ def stack_arrays(
         with self.device_scope(device_id, framework):
             return self._operations.stack(arrays, framework)
 
+    def normalize_planes(
+        self,
+        data: Any,
+        dtype: Any,
+        scales: Sequence[float | None],
+        module: Any | None = None,
+    ) -> Any:
+        """Normalize one owned plane stack through its native operation owner."""
+        if len(scales) != len(data):
+            raise ValueError("Normalization scales must match the leading plane axis.")
+        framework = module if module is not None else self.import_module()
+        with self._array_device_scope(data, framework):
+            return self._operations.normalize_planes(data, dtype, scales, framework)
+
     def scale_dtype(
         self,
         data: Any,
diff --git a/tests/test_native_geometry.py b/tests/test_native_geometry.py
index 8bcfc17..8e9a9cd 100644
--- a/tests/test_native_geometry.py
+++ b/tests/test_native_geometry.py
@@ -63,3 +63,18 @@ def __array__(self):
         PYCLESPERANTO_OPERATIONS.broadcast_to(source, (4, 2, 3), None)
     with pytest.raises(NotImplementedError, match="dtype"):
         PYCLESPERANTO_OPERATIONS.ones_like(source, (2, 3), bool, None)
+
+
+@pytest.mark.parametrize("dtype", (np.float32, np.complex64, np.int32))
+def test_normalize_planes_preserves_source_and_division_dtype(dtype):
+    source = np.arange(24, dtype=dtype).reshape(2, 3, 4)[:, :, ::-1]
+    original = source.copy()
+    expected = np.stack((source[0] / 3.0, source[1]))
+    actual = MemoryType.NUMPY.normalize_planes(source, dtype, (3.0, None))
+    np.testing.assert_array_equal(actual, expected)
+    np.testing.assert_array_equal(source, original)
+    assert actual.dtype == expected.dtype
+    assert not np.shares_memory(actual, source)
+    for scales in ((3.0,), (3.0, None, 2.0)):
+        with pytest.raises(ValueError, match="leading plane axis"):
+            MemoryType.NUMPY.normalize_planes(source, dtype, scales)