Add device-neutral array geometry

Author: trissim

docs/source/api_reference.rst

@@ -12,6 +12,11 @@ Public imports
 ``detect_memory_type`` and ``convert_memory``
   Detection and explicit conversion.
 
+``ArrayGeometry``
+  Framework-neutral shape inspection for arrays and nominal ``ArrayPayload``
+  values. Geometry inspection reads declared shape metadata and does not move
+  device data to host memory.
+
 ``memory_types``, ``numpy``, ``cupy``, ``torch``, ``tensorflow``, ``jax``, ``pyclesperanto``
   Callable memory declarations and wrappers.
 

docs/source/conf.py

@@ -22,7 +22,7 @@
 project = 'arraybridge'
 copyright = '2025, Tristan Simas'
 author = 'Tristan Simas'
-release = '0.3.0'
+release = '0.3.1'
 
 # -- General configuration ---------------------------------------------------
 # https://www.sphinx-doc.org/en/master/usage/configuration.html#general-configuration

pyproject.toml

@@ -5,7 +5,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "arraybridge"
-version = "0.3.0"
+version = "0.3.1"
 description = "Unified API for NumPy, CuPy, PyTorch, TensorFlow, JAX, and pyclesperanto with automatic memory type conversion"
 authors = [{name = "Tristan Simas", email = "tristan.simas@mail.mcgill.ca"}]
 license = {text = "MIT"}

src/arraybridge/__init__.py

@@ -5,9 +5,10 @@
 and unified utilities for working with multiple array/tensor frameworks.
 """
 
-__version__ = "0.3.0"
+__version__ = "0.3.1"
 
 from . import decorators as _decorators
+from .array_geometry import ArrayGeometry
 from .array_payload import ArrayPayload
 from .converters import convert_memory, detect_memory_type
 from .dtype_scaling import SCALING_FUNCTIONS
@@ -39,6 +40,7 @@
     "MemoryType",
     "MemoryContractAttribute",
     "ArrayPayload",
+    "ArrayGeometry",
     "CPU_MEMORY_TYPES",
     "GPU_MEMORY_TYPES",
     "SUPPORTED_MEMORY_TYPES",

src/arraybridge/array_geometry.py

@@ -0,0 +1,57 @@
+"""Framework-neutral array geometry inspection."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+import numpy as np
+
+from arraybridge.array_payload import ArrayPayload
+
+
+@dataclass(frozen=True, slots=True)
+class ArrayGeometry:
+    """Concrete shape metadata without moving array data between frameworks."""
+
+    shape: tuple[int, ...]
+
+    @property
+    def ndim(self) -> int:
+        """Return the rank derived from the canonical shape."""
+
+        return len(self.shape)
+
+    @classmethod
+    def from_value(cls, value: Any) -> ArrayGeometry | None:
+        """Inspect an array or nominal payload without forcing host conversion."""
+
+        data = value.array_payload_data() if isinstance(value, ArrayPayload) else value
+        declared_shape = getattr(data, "shape", None)
+        if declared_shape is not None:
+            try:
+                return cls(tuple(int(axis_size) for axis_size in declared_shape))
+            except (TypeError, ValueError):
+                return None
+
+        try:
+            array = np.asarray(data)
+        except (TypeError, ValueError):
+            return None
+        return cls(tuple(int(axis_size) for axis_size in array.shape))
+
+    @classmethod
+    def require_from_value(
+        cls,
+        value: Any,
+        *,
+        value_name: str = "Value",
+    ) -> ArrayGeometry:
+        """Return concrete geometry or reject a value without an array shape."""
+
+        geometry = cls.from_value(value)
+        if geometry is None:
+            raise TypeError(
+                f"{value_name} requires concrete array geometry, got " f"{type(value).__name__}."
+            )
+        return geometry

tests/test_array_geometry.py

@@ -0,0 +1,53 @@
+from __future__ import annotations
+
+import numpy as np
+import pytest
+
+from arraybridge import ArrayGeometry, ArrayPayload
+
+
+class DeviceArray:
+    """Array-shaped value that rejects implicit host conversion."""
+
+    shape = (2, 3, 4)
+
+    def __array__(self, dtype=None):
+        del dtype
+        raise TypeError("implicit host conversion is forbidden")
+
+
+class DevicePayload(ArrayPayload):
+    def __init__(self, data):
+        self.data = data
+
+    def array_payload_data(self):
+        return self.data
+
+    def with_data(self, data):
+        return type(self)(data)
+
+
+def test_array_geometry_reads_declared_device_shape_without_host_conversion() -> None:
+    assert ArrayGeometry.from_value(DeviceArray()) == ArrayGeometry((2, 3, 4))
+
+
+def test_array_geometry_unwraps_nominal_array_payload() -> None:
+    assert ArrayGeometry.from_value(DevicePayload(DeviceArray())) == ArrayGeometry((2, 3, 4))
+
+
+def test_array_geometry_falls_back_for_python_array_inputs() -> None:
+    assert ArrayGeometry.from_value([[1, 2], [3, 4]]) == ArrayGeometry((2, 2))
+    assert ArrayGeometry.from_value(np.zeros((4, 5))) == ArrayGeometry((4, 5))
+
+
+def test_array_geometry_requires_concrete_shape() -> None:
+    class UnshapedValue:
+        def __array__(self, dtype=None):
+            del dtype
+            raise TypeError("not an array")
+
+    with pytest.raises(TypeError, match="Runtime output requires concrete array geometry"):
+        ArrayGeometry.require_from_value(
+            UnshapedValue(),
+            value_name="Runtime output",
+        )