From 6cb7859b47c8cad1f584c4ce0e2f07a799dfbed6 Mon Sep 17 00:00:00 2001 From: nicola-bastianello Date: Fri, 12 Jun 2026 16:28:06 +0200 Subject: [PATCH 01/23] feat: add dtype objects, backend properties, and dtype registration --- .../interoperability/_abstracts/backend.py | 132 +++++++++++++++ .../interoperability/_jax/jax_backend.py | 79 ++++++++- .../interoperability/_numpy/numpy_backend.py | 84 +++++++++ .../_pytorch/pytorch_backend.py | 62 +++++++ .../_tensorflow/tensorflow_backend.py | 82 +++++++++ decent_array/types.py | 160 ++++++++++++++++++ 6 files changed, 598 insertions(+), 1 deletion(-) diff --git a/decent_array/interoperability/_abstracts/backend.py b/decent_array/interoperability/_abstracts/backend.py index cfa17b8..90d8759 100644 --- a/decent_array/interoperability/_abstracts/backend.py +++ b/decent_array/interoperability/_abstracts/backend.py @@ -325,3 +325,135 @@ def uniform_like(self, x: Array, low: float = 0.0, high: float = 1.0) -> Array: @abstractmethod def choice(self, x: Array, size: int, replace: bool = True) -> Array: """Sample ``size`` elements from ``x``.""" + + # DTYPES -------------------------------------------------------------- + + @property + @abstractmethod + def bool(self) -> Any: # noqa: ANN401 + """Bool dtype.""" + + @property + @abstractmethod + def uint8(self) -> Any: # noqa: ANN401 + """Unsigned 8-bit integer dtype.""" + + @property + def uint16(self) -> Any | None: # noqa: ANN401 + """Unsigned 16-bit integer dtype (optional).""" + return None + + @property + def uint32(self) -> Any | None: # noqa: ANN401 + """Unsigned 32-bit integer dtype (optional).""" + return None + + @property + def uint64(self) -> Any | None: # noqa: ANN401 + """Unsigned 64-bit integer dtype (optional).""" + return None + + @property + @abstractmethod + def int8(self) -> Any: # noqa: ANN401 + """Signed 8-bit integer dtype.""" + + @property + @abstractmethod + def int16(self) -> Any: # noqa: ANN401 + """Signed 16-bit integer dtype.""" + + @property + @abstractmethod + def int32(self) -> Any: # noqa: ANN401 + """Signed 32-bit integer dtype.""" + + @property + @abstractmethod + def int64(self) -> Any: # noqa: ANN401 + """Signed 64-bit integer dtype.""" + + @property + @abstractmethod + def float16(self) -> Any: # noqa: ANN401 + """16-bit floating-point dtype.""" + + @property + @abstractmethod + def float32(self) -> Any: # noqa: ANN401 + """32-bit floating-point dtype.""" + + @property + @abstractmethod + def float64(self) -> Any: # noqa: ANN401 + """64-bit floating-point dtype.""" + + @property + @abstractmethod + def complex64(self) -> Any: # noqa: ANN401 + """64-bit complex dtype.""" + + @property + @abstractmethod + def complex128(self) -> Any: # noqa: ANN401 + """128-bit complex dtype.""" + + @property + def float128(self) -> Any | None: # noqa: ANN401 + """128-bit floating-point dtype (optional).""" + return None + + @property + def complex256(self) -> Any | None: # noqa: ANN401 + """256-bit complex dtype (optional).""" + return None + + @property + def qint8(self) -> Any | None: # noqa: ANN401 + """Quantized 8-bit integer dtype (optional).""" + return None + + @property + def qint16(self) -> Any | None: # noqa: ANN401 + """Quantized 16-bit integer dtype (optional).""" + return None + + @property + def qint32(self) -> Any | None: # noqa: ANN401 + """Quantized 32-bit integer dtype (optional).""" + return None + + @property + def quint8(self) -> Any | None: # noqa: ANN401 + """Quantized unsigned 8-bit integer dtype (optional).""" + return None + + @property + def quint16(self) -> Any | None: # noqa: ANN401 + """Quantized unsigned 16-bit integer dtype (optional).""" + return None + + @property + def bfloat16(self) -> Any | None: # noqa: ANN401 + """Brain floating point 16-bit dtype (optional).""" + return None + + @property + def unicode(self) -> Any | None: # noqa: ANN401 + """Unicode string dtype (optional).""" + return None + + @property + def bytes(self) -> Any | None: # noqa: ANN401 + """Byte string dtype (optional).""" + return None + + @property + def object(self) -> Any | None: # noqa: ANN401 + """Python object dtype (optional).""" + return None + + @property + def void(self) -> Any | None: # noqa: ANN401 + """Raw/void dtype (optional).""" + return None diff --git a/decent_array/interoperability/_jax/jax_backend.py b/decent_array/interoperability/_jax/jax_backend.py index 734b634..6a10fec 100644 --- a/decent_array/interoperability/_jax/jax_backend.py +++ b/decent_array/interoperability/_jax/jax_backend.py @@ -10,7 +10,7 @@ from __future__ import annotations -from collections.abc import Sequence +from collections.abc import Callable, Sequence from time import time_ns from typing import Any, cast @@ -318,5 +318,82 @@ def _next_key(self) -> jax.Array: self._key, sub = jax.random.split(self._key) return cast("jax.Array", sub) + # Dtypes + + @property + def bool(self) -> Any: # noqa: ANN401 + return np.dtype(jnp.bool_) + + @property + def uint8(self) -> Any: # noqa: ANN401 + return np.dtype(jnp.uint8) + + @property + def uint16(self) -> Any: # noqa: ANN401 + return np.dtype(jnp.uint16) + + @property + def uint32(self) -> Any: # noqa: ANN401 + return np.dtype(jnp.uint32) + + @property + def uint64(self) -> Any: # noqa: ANN401 + if _x64_enabled(): + return np.dtype(jnp.uint64) + return None + + @property + def int8(self) -> Any: # noqa: ANN401 + return np.dtype(jnp.int8) + + @property + def int16(self) -> Any: # noqa: ANN401 + return np.dtype(jnp.int16) + + @property + def int32(self) -> Any: # noqa: ANN401 + return np.dtype(jnp.int32) + + @property + def int64(self) -> Any: # noqa: ANN401 + if _x64_enabled(): + return np.dtype(jnp.int64) + return None + + @property + def float16(self) -> Any: # noqa: ANN401 + return np.dtype(jnp.float16) + + @property + def float32(self) -> Any: # noqa: ANN401 + return np.dtype(jnp.float32) + + @property + def float64(self) -> Any: # noqa: ANN401 + if _x64_enabled(): + return np.dtype(jnp.float64) + return None + + @property + def complex64(self) -> Any: # noqa: ANN401 + return np.dtype(jnp.complex64) + + @property + def complex128(self) -> Any: # noqa: ANN401 + if _x64_enabled(): + return np.dtype(jnp.complex128) + return None + + @property + def bfloat16(self) -> Any: # noqa: ANN401 + return np.dtype(jnp.bfloat16) + + +_JAX_CONFIG_READ = cast("Callable[[str], Any]", jax.config.read) + + +def _x64_enabled() -> bool: + return bool(_JAX_CONFIG_READ("jax_enable_x64")) + register_backend(SupportedFrameworks.JAX, JaxBackend) diff --git a/decent_array/interoperability/_numpy/numpy_backend.py b/decent_array/interoperability/_numpy/numpy_backend.py index 72f363d..12a9358 100644 --- a/decent_array/interoperability/_numpy/numpy_backend.py +++ b/decent_array/interoperability/_numpy/numpy_backend.py @@ -312,5 +312,89 @@ def uniform_like(self, x: Array, low: float = 0.0, high: float = 1.0) -> Array: def choice(self, x: Array, size: int, replace: bool = True) -> Array: return Array(self._rng.choice(x.value, size=size, replace=replace)) + # Dtypes + + @property + def bool(self) -> Any: # noqa: ANN401 + return np.dtype(np.bool_) + + @property + def uint8(self) -> Any: # noqa: ANN401 + return np.dtype(np.uint8) + + @property + def uint16(self) -> Any: # noqa: ANN401 + return np.dtype(np.uint16) + + @property + def uint32(self) -> Any: # noqa: ANN401 + return np.dtype(np.uint32) + + @property + def uint64(self) -> Any: # noqa: ANN401 + return np.dtype(np.uint64) + + @property + def int8(self) -> Any: # noqa: ANN401 + return np.dtype(np.int8) + + @property + def int16(self) -> Any: # noqa: ANN401 + return np.dtype(np.int16) + + @property + def int32(self) -> Any: # noqa: ANN401 + return np.dtype(np.int32) + + @property + def int64(self) -> Any: # noqa: ANN401 + return np.dtype(np.int64) + + @property + def float16(self) -> Any: # noqa: ANN401 + return np.dtype(np.float16) + + @property + def float32(self) -> Any: # noqa: ANN401 + return np.dtype(np.float32) + + @property + def float64(self) -> Any: # noqa: ANN401 + return np.dtype(np.float64) + + @property + def complex64(self) -> Any: # noqa: ANN401 + return np.dtype(np.complex64) + + @property + def complex128(self) -> Any: # noqa: ANN401 + return np.dtype(np.complex128) + + @property + def float128(self) -> Any | None: # noqa: ANN401 + float128 = getattr(np, "float128", None) + return np.dtype(float128) if float128 is not None else None + + @property + def complex256(self) -> Any | None: # noqa: ANN401 + complex256 = getattr(np, "complex256", None) + return np.dtype(complex256) if complex256 is not None else None + + @property + def unicode(self) -> Any: # noqa: ANN401 + return np.dtype(np.str_) + + @property + def bytes(self) -> Any: # noqa: ANN401 + return np.dtype(np.bytes_) + + @property + def object(self) -> Any: # noqa: ANN401 + return np.dtype(np.object_) + + @property + def void(self) -> Any: # noqa: ANN401 + return np.dtype(np.void) + register_backend(SupportedFrameworks.NUMPY, NumpyBackend) diff --git a/decent_array/interoperability/_pytorch/pytorch_backend.py b/decent_array/interoperability/_pytorch/pytorch_backend.py index 1439f4f..92494fd 100644 --- a/decent_array/interoperability/_pytorch/pytorch_backend.py +++ b/decent_array/interoperability/_pytorch/pytorch_backend.py @@ -359,5 +359,67 @@ def choice(self, x: Array, size: int, replace: bool = True) -> Array: indices = weights.multinomial(num_samples=size, replacement=replace, generator=self._generator) return Array(v[indices]) + # Dtypes + + @property + def bool(self) -> torch.dtype: + return torch.bool + + @property + def uint8(self) -> torch.dtype: + return torch.uint8 + + @property + def int8(self) -> torch.dtype: + return torch.int8 + + @property + def int16(self) -> torch.dtype: + return torch.int16 + + @property + def int32(self) -> torch.dtype: + return torch.int32 + + @property + def int64(self) -> torch.dtype: + return torch.int64 + + @property + def float16(self) -> torch.dtype: + return torch.float16 + + @property + def float32(self) -> torch.dtype: + return torch.float32 + + @property + def float64(self) -> torch.dtype: + return torch.float64 + + @property + def complex64(self) -> torch.dtype: + return torch.complex64 + + @property + def complex128(self) -> torch.dtype: + return torch.complex128 + + @property + def qint8(self) -> Any: # noqa: ANN401 + return torch.qint8 + + @property + def qint32(self) -> Any: # noqa: ANN401 + return torch.qint32 + + @property + def quint8(self) -> Any: # noqa: ANN401 + return torch.quint8 + + @property + def bfloat16(self) -> torch.dtype: + return torch.bfloat16 + register_backend(SupportedFrameworks.PYTORCH, PyTorchBackend) diff --git a/decent_array/interoperability/_tensorflow/tensorflow_backend.py b/decent_array/interoperability/_tensorflow/tensorflow_backend.py index 586acae..ae3bcaa 100644 --- a/decent_array/interoperability/_tensorflow/tensorflow_backend.py +++ b/decent_array/interoperability/_tensorflow/tensorflow_backend.py @@ -366,5 +366,87 @@ def choice(self, x: Array, size: int, replace: bool = True) -> Array: indices = tf.cast(tf.math.top_k(scores, k=size).indices, tf.int32) return Array(tf.gather(v, indices)) + # Dtypes + + @property + def bool(self) -> tf.dtypes.DType: + return tf.bool + + @property + def uint8(self) -> tf.dtypes.DType: + return tf.uint8 + + @property + def uint16(self) -> tf.dtypes.DType: + return tf.uint16 + + @property + def uint32(self) -> tf.dtypes.DType: + return tf.uint32 + + @property + def uint64(self) -> tf.dtypes.DType: + return tf.uint64 + + @property + def int8(self) -> tf.dtypes.DType: + return tf.int8 + + @property + def int16(self) -> tf.dtypes.DType: + return tf.int16 + + @property + def int32(self) -> tf.dtypes.DType: + return tf.int32 + + @property + def int64(self) -> tf.dtypes.DType: + return tf.int64 + + @property + def float16(self) -> tf.dtypes.DType: + return tf.float16 + + @property + def float32(self) -> tf.dtypes.DType: + return tf.float32 + + @property + def float64(self) -> tf.dtypes.DType: + return tf.float64 + + @property + def complex64(self) -> tf.dtypes.DType: + return tf.complex64 + + @property + def complex128(self) -> tf.dtypes.DType: + return tf.complex128 + + @property + def qint8(self) -> tf.dtypes.DType: + return tf.qint8 + + @property + def qint16(self) -> tf.dtypes.DType: + return tf.qint16 + + @property + def qint32(self) -> tf.dtypes.DType: + return tf.qint32 + + @property + def quint8(self) -> tf.dtypes.DType: + return tf.quint8 + + @property + def quint16(self) -> tf.dtypes.DType: + return tf.quint16 + + @property + def bfloat16(self) -> tf.dtypes.DType: + return tf.bfloat16 + register_backend(SupportedFrameworks.TENSORFLOW, TensorflowBackend) diff --git a/decent_array/types.py b/decent_array/types.py index 7e39af0..71d4ded 100644 --- a/decent_array/types.py +++ b/decent_array/types.py @@ -2,9 +2,12 @@ from __future__ import annotations +import sys from enum import Enum from typing import TYPE_CHECKING, SupportsIndex, TypeAlias, Union +from decent_array.interoperability._backend_manager import register_backend_listener + if TYPE_CHECKING: import jax import numpy @@ -12,6 +15,35 @@ import torch from decent_array._array import Array + from decent_array.interoperability._abstracts import Backend + + +_BACKEND_INSTANCE: Backend | None = None +_error = RuntimeError("No backend active: call 'set_backend' with a supported framework to activate one.") + + +def _update_backend(backend: Backend | None) -> None: + global _BACKEND_INSTANCE # noqa: PLW0603 + _BACKEND_INSTANCE = backend + + _reset_dtype_registry() + + if backend is not None: + _initialize_dtype_registry(backend) + + +def _reset_dtype_registry() -> None: + """Remove dtype module attributes and clear the native reverse map.""" + module = sys.modules[__name__] + + _NATIVE_TO_IOP.clear() + for name in _SUPPORTED_DTYPES + _OPTIONAL_DTYPES: + if hasattr(module, name): + delattr(module, name) + + +register_backend_listener(_update_backend) + ArrayLike: TypeAlias = Union["numpy.ndarray", "torch.Tensor", "tf.Tensor", "jax.Array"] # noqa: UP040 """ @@ -70,3 +102,131 @@ class DTypes(Enum): FLOAT64 = "float64" COMPLEX64 = "complex64" COMPLEX128 = "complex128" + + +class dtype: # noqa: N801 + """Base class for dtypes.""" + + def __init__(self, name: str): + if _BACKEND_INSTANCE is None: + raise _error + + # name doesn't map to any supported/optional dtype + if name not in (_SUPPORTED_DTYPES + _OPTIONAL_DTYPES): + raise ValueError(f"dtype {name} is not supported.") + + # native_dtype is None if the dtype is not supported + native_dtype = getattr(_BACKEND_INSTANCE, name, None) + if native_dtype is None: + raise ValueError(f"dtype {name} is not supported.") # TODO add device and backend name to contextualize + + self._name = name + self._native_dtype = native_dtype # TODO better name? + + @property + def name(self) -> str: + """Name of the dtype.""" + return self._name + + def __str__(self) -> str: + """Name of the dtype.""" + return self.name + + def __eq__(self, other: object) -> bool: + """Check equivalence by ``name`` attributes.""" + if not isinstance(other, dtype): + return NotImplemented + return self.name == other.name + + def __hash__(self) -> int: + """Hash of the dtype.""" + return hash(self.name) + + +_SUPPORTED_DTYPES = [ + "bool", + "uint8", + "int8", + "int16", + "int32", + "int64", + "float16", + "float32", + "float64", + "complex64", + "complex128", +] + +_OPTIONAL_DTYPES = [ + "uint16", + "uint32", + "uint64", + "float128", + "complex256", + "qint8", + "qint16", + "qint32", + "quint8", + "quint16", + "bfloat16", + "unicode", + "bytes", + "object", + "void", +] + + +# TODO jax needs jax.config.x64_enabled=True to use 64-bit; if not, it will silently downcast to 32-bit +# right now I return None if jax.config.x64_enabled=False, because it is effectively unavailable. do we want to do this +# or let it silently downcast? + +# TODO do we want to have an optional dependence on ml_dtypes to support bfloat16 for numpy? + +# ✓: available, ~: limited support, ✗: not available +# dtype | NumPy | JAX | PyTorch | TensorFlow | Notes # noqa: ERA001 +# ---------------------------|-------|-----------|------------|------------|------------------------------------------ +# uint16 | ✓ | ✓ | ~ | ✓ | https://github.com/pytorch/pytorch/issues/58734 +# uint32 | ✓ | ✓ | ~ | ✓ | https://github.com/pytorch/pytorch/issues/58734 +# uint64 | ✓ | ✓ | ~ | ✓ | https://github.com/pytorch/pytorch/issues/58734 + +# float128 | ✓ | ✗ | ✗ | ✗ | platform-dependent +# complex256 | ✓ | ✗ | ✗ | ✗ | platform-dependent + +# qint8 | ✗ | ✗ | ✓ | ✓ | +# qint16 | ✗ | ✗ | ✗ | ✓ | +# qint32 | ✗ | ✗ | ✓ | ✓ | + +# quint8 | ✗ | ✗ | ✓ | ✓ | +# quint16 | ✗ | ✗ | ✗ | ✓ | + +# bfloat16 | ✗* | ✓ | ✓ | ✓ | * available with ml_dtypes dependency + +# unicode_ | ✓ | ✗ | ✗ | ✗ | np.str_ +# bytes_ | ✓ | ✗ | ✗ | ✗ | np.bytes_, tf.string +# object | ✓ | ✗ | ✗ | ✗ | +# void | ✓ | ✗ | ✗ | ✗ | + + +# dict for reverse search native dtype -> iop dtype +_NATIVE_TO_IOP: dict[object, dtype] = {} + + +def _initialize_dtype_registry(backend: Backend) -> None: + """Initialize dtype objects and the reverse native-to-iop map for ``backend``.""" + global _NATIVE_TO_IOP # noqa: PLW0603 + module = sys.modules[__name__] + _NATIVE_TO_IOP = {} + + for name in (_SUPPORTED_DTYPES + _OPTIONAL_DTYPES): + _register_dtype(module, backend, name) + + +def _register_dtype(module: object, backend: Backend, name: str) -> None: + """Create and publish a dtype (if available).""" + native_dtype = getattr(backend, name, None) + if native_dtype is None: + return + + iop_dtype = dtype(name) + setattr(module, name, iop_dtype) + _NATIVE_TO_IOP[native_dtype] = iop_dtype From 98236971df601e898ce20e4b2dd1865e5710ee4c Mon Sep 17 00:00:00 2001 From: nicola-bastianello Date: Fri, 12 Jun 2026 20:31:05 +0200 Subject: [PATCH 02/23] docs: added discussion of supported dtypes --- decent_array/types.py | 2 +- docs/source/user.rst | 248 ++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 249 insertions(+), 1 deletion(-) diff --git a/decent_array/types.py b/decent_array/types.py index 71d4ded..8585339 100644 --- a/decent_array/types.py +++ b/decent_array/types.py @@ -202,7 +202,7 @@ def __hash__(self) -> int: # bfloat16 | ✗* | ✓ | ✓ | ✓ | * available with ml_dtypes dependency # unicode_ | ✓ | ✗ | ✗ | ✗ | np.str_ -# bytes_ | ✓ | ✗ | ✗ | ✗ | np.bytes_, tf.string +# bytes_ | ✓ | ✗ | ✗ | ✓ | np.bytes_, tf.string # object | ✓ | ✗ | ✗ | ✗ | # void | ✓ | ✗ | ✗ | ✗ | diff --git a/docs/source/user.rst b/docs/source/user.rst index 9c08874..ad409ba 100644 --- a/docs/source/user.rst +++ b/docs/source/user.rst @@ -8,3 +8,251 @@ Requires `Python 3.13+ `_ .. code-block:: bash pip install decent-array + + +dtypes +------ + +decent-array exposes a number of dtypes which are bound to the corresponding framework-native dtypes. The following +lists collect the dtypes always supported, and the dtypes supported only by some frameworks/under certain conditions. +The table below reports all the details. + + + +dtypes always available +~~~~~~~~~~~~~~~~~~~~~~~ + +- bool +- int8 +- int16 +- int32 +- uint8 +- float16 +- float32 +- complex64 + +dtypes conditionally available +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +- int64 +- uint16 +- uint32 +- uint64 +- bfloat16 +- float64 +- float128 +- complex128 +- complex256 +- qint8 +- quint8 +- qint16 +- quint16 +- qint32 +- unicode +- bytes +- object +- void + + +.. list-table:: dtype support across frameworks + :header-rows: 1 + :widths: 22 10 10 10 12 36 + + * - dtype + - NumPy + - JAX + - PyTorch + - TensorFlow + - Notes + * - ``bool_`` + - ✓ + - ✓ + - ✓ + - ✓ + - + * - **Integers** + - + - + - + - + - + * - ``int8`` + - ✓ + - ✓ + - ✓ + - ✓ + - + * - ``int16`` + - ✓ + - ✓ + - ✓ + - ✓ + - + * - ``int32`` + - ✓ + - ✓ + - ✓ + - ✓ + - + * - ``int64`` + - ✓ + - ⚠️ + - ✓ + - ✓ + - JAX requires ``jax_enable_x64=True`` + * - **Unsigned integers** + - + - + - + - + - + * - ``uint8`` + - ✓ + - ✓ + - ✓ + - ✓ + - + * - ``uint16`` + - ✓ + - ⚠️ + - ⚠️ + - ✓ + - JAX requires ``jax_enable_x64=True``; PyTorch support is limited/experimental + * - ``uint32`` + - ✓ + - ⚠️ + - ⚠️ + - ✓ + - JAX requires ``jax_enable_x64=True``; PyTorch support is limited/experimental + * - ``uint64`` + - ✓ + - ⚠️ + - ⚠️ + - ✓ + - JAX requires ``jax_enable_x64=True``; PyTorch support is limited/experimental + * - **Floating point** + - + - + - + - + - + * - ``float16`` + - ✓ + - ✓ + - ✓ + - ✓ + - + * - ``bfloat16`` + - ✗ + - ✓ + - ✓ + - ✓ + - + * - ``float32`` + - ✓ + - ✓ + - ✓ + - ✓ + - + * - ``float64`` + - ✓ + - ⚠️ + - ✓ + - ✓ + - JAX requires ``jax_enable_x64=True`` + * - ``float128`` + - ⚠️ + - ✗ + - ✗ + - ✗ + - NumPy support is platform-dependent + * - **Complex** + - + - + - + - + - + * - ``complex64`` + - ✓ + - ✓ + - ✓ + - ✓ + - + * - ``complex128`` + - ✓ + - ⚠️ + - ✓ + - ✓ + - JAX requires ``jax_enable_x64=True`` + * - ``complex256`` + - ⚠️ + - ✗ + - ✗ + - ✗ + - NumPy support is platform-dependent + * - **Quantized** + - + - + - + - + - + * - ``qint8`` + - ✗ + - ✗ + - ✓ + - ✓ + - + * - ``quint8`` + - ✗ + - ✗ + - ✓ + - ✓ + - + * - ``qint16`` + - ✗ + - ✗ + - ✗ + - ✓ + - + * - ``quint16`` + - ✗ + - ✗ + - ✗ + - ✓ + - + * - ``qint32`` + - ✗ + - ✗ + - ✓ + - ✓ + - + * - **Miscellaneous** + - + - + - + - + - + * - ``unicode_`` + - ✓ + - ✗ + - ✗ + - ✗ + - Equivalent to ``np.str_`` + * - ``bytes_`` + - ✓ + - ✗ + - ✗ + - ✓ + - Equivalent to ``np.bytes_`` and ``tf.string`` + * - ``object`` + - ✓ + - ✗ + - ✗ + - ✗ + - + * - ``void`` + - ✓ + - ✗ + - ✗ + - ✗ + - \ No newline at end of file From 1c18f8af9ba2216a76375cee064edd727e9efd33 Mon Sep 17 00:00:00 2001 From: nicola-bastianello Date: Sun, 14 Jun 2026 16:27:20 +0200 Subject: [PATCH 03/23] enh: next iteration of dtypes --- decent_array/interoperability/__init__.py | 2 +- .../_iop/{comparasion.py => comparison.py} | 0 decent_array/types.py | 235 ------------------ decent_array/types/__init__.py | 51 ++++ decent_array/types/_dtypes.py | 141 +++++++++++ decent_array/types/_types.py | 91 +++++++ 6 files changed, 284 insertions(+), 236 deletions(-) rename decent_array/interoperability/_iop/{comparasion.py => comparison.py} (100%) delete mode 100644 decent_array/types.py create mode 100644 decent_array/types/__init__.py create mode 100644 decent_array/types/_dtypes.py create mode 100644 decent_array/types/_types.py diff --git a/decent_array/interoperability/__init__.py b/decent_array/interoperability/__init__.py index 8465bca..2fcdd50 100644 --- a/decent_array/interoperability/__init__.py +++ b/decent_array/interoperability/__init__.py @@ -21,7 +21,7 @@ bitwise_right_shift, bitwise_xor, ) -from ._iop.comparasion import equal, greater, greater_equal, less, less_equal, not_equal +from ._iop.comparison import equal, greater, greater_equal, less, less_equal, not_equal from ._iop.creation import eye, ones, ones_like, zeros, zeros_like from ._iop.linalg import dot, matmul, norm, vecdot, vector_norm from ._iop.manipulations import ( diff --git a/decent_array/interoperability/_iop/comparasion.py b/decent_array/interoperability/_iop/comparison.py similarity index 100% rename from decent_array/interoperability/_iop/comparasion.py rename to decent_array/interoperability/_iop/comparison.py diff --git a/decent_array/types.py b/decent_array/types.py deleted file mode 100644 index 4fa2b07..0000000 --- a/decent_array/types.py +++ /dev/null @@ -1,235 +0,0 @@ -"""Type definitions for optimization variables.""" - -from __future__ import annotations - -import sys -from enum import Enum -from typing import TYPE_CHECKING, SupportsIndex, TypeAlias, Union - -from decent_array.interoperability._backend_manager import register_backend_listener - -if TYPE_CHECKING: - import jax - import numpy - import tensorflow as tf - import torch - - from decent_array._array import Array - from decent_array.interoperability._abstracts import Backend - - -_BACKEND_INSTANCE: Backend | None = None -_error = RuntimeError("No backend active: call 'set_backend' with a supported framework to activate one.") - - -def _update_backend(backend: Backend | None) -> None: - global _BACKEND_INSTANCE # noqa: PLW0603 - _BACKEND_INSTANCE = backend - - _reset_dtype_registry() - - if backend is not None: - _initialize_dtype_registry(backend) - - -def _reset_dtype_registry() -> None: - """Remove dtype module attributes and clear the native reverse map.""" - module = sys.modules[__name__] - - _NATIVE_TO_IOP.clear() - for name in _SUPPORTED_DTYPES + _OPTIONAL_DTYPES: - if hasattr(module, name): - delattr(module, name) - - -register_backend_listener(_update_backend) - - -ArrayLike: TypeAlias = Union["numpy.ndarray", "torch.Tensor", "tf.Tensor", "jax.Array"] # noqa: UP040 -""" -Type alias for array-like types supported in decent-array, including NumPy arrays, -PyTorch tensors, TensorFlow tensors, and JAX arrays. -""" - -SupportedArrayTypes: TypeAlias = bool | int | float | complex | ArrayLike # noqa: UP040 -""" -Type alias for supported types for optimization variables in decent-array, -including array-like types and scalars. -""" - -ArrayKey: TypeAlias = ( # noqa: UP040 - "int | SupportsIndex | slice | Array | tuple[int | SupportsIndex | slice | Array | None, ...] | None" -) -""" -Type alias for valid keys used to index into supported array types. -Includes single indices, tuples of indices, slices, and tuples of slices. -""" - - -# Its important that the enum values correspond to the folder names of the backends, -# since those are used for dynamic imports in _backend_manager.py -class SupportedFrameworks(Enum): - """Enum for supported frameworks in decent-array.""" - - NUMPY = "numpy" - PYTORCH = "pytorch" - TENSORFLOW = "tensorflow" - JAX = "jax" - - -class SupportedDevices(Enum): - """Enum for supported devices in decent-array.""" - - CPU = "cpu" - GPU = "gpu" - MPS = "mps" - - -class DTypes(Enum): - """Enum for supported dtypes in decent-array.""" - - BOOL = "bool" - UINT8 = "uint8" - UINT16 = "uint16" - UINT32 = "uint32" - UINT64 = "uint64" - INT8 = "int8" - INT16 = "int16" - INT32 = "int32" - INT64 = "int64" - FLOAT16 = "float16" - FLOAT32 = "float32" - FLOAT64 = "float64" - COMPLEX64 = "complex64" - COMPLEX128 = "complex128" - - -_STRING_TO_DTYPE = {dt.value: dt for dt in DTypes} - - -class dtype: # noqa: N801 - """Base class for dtypes.""" - - def __init__(self, name: str): - if _BACKEND_INSTANCE is None: - raise _error - - # name doesn't map to any supported/optional dtype - if name not in (_SUPPORTED_DTYPES + _OPTIONAL_DTYPES): - raise ValueError(f"dtype {name} is not supported.") - - # native_dtype is None if the dtype is not supported - native_dtype = getattr(_BACKEND_INSTANCE, name, None) - if native_dtype is None: - raise ValueError(f"dtype {name} is not supported.") # TODO add device and backend name to contextualize - - self._name = name - self._native_dtype = native_dtype # TODO better name? - - @property - def name(self) -> str: - """Name of the dtype.""" - return self._name - - def __str__(self) -> str: - """Name of the dtype.""" - return self.name - - def __eq__(self, other: object) -> bool: - """Check equivalence by ``name`` attributes.""" - if not isinstance(other, dtype): - return NotImplemented - return self.name == other.name - - def __hash__(self) -> int: - """Hash of the dtype.""" - return hash(self.name) - - -_SUPPORTED_DTYPES = [ - "bool", - "uint8", - "int8", - "int16", - "int32", - "int64", - "float16", - "float32", - "float64", - "complex64", - "complex128", -] - -_OPTIONAL_DTYPES = [ - "uint16", - "uint32", - "uint64", - "float128", - "complex256", - "qint8", - "qint16", - "qint32", - "quint8", - "quint16", - "bfloat16", - "unicode", - "bytes", - "object", - "void", -] - - -# TODO jax needs jax.config.x64_enabled=True to use 64-bit; if not, it will silently downcast to 32-bit -# right now I return None if jax.config.x64_enabled=False, because it is effectively unavailable. do we want to do this -# or let it silently downcast? - -# TODO do we want to have an optional dependence on ml_dtypes to support bfloat16 for numpy? - -# ✓: available, ~: limited support, ✗: not available -# dtype | NumPy | JAX | PyTorch | TensorFlow | Notes # noqa: ERA001 -# ---------------------------|-------|-----------|------------|------------|------------------------------------------ -# uint16 | ✓ | ✓ | ~ | ✓ | https://github.com/pytorch/pytorch/issues/58734 -# uint32 | ✓ | ✓ | ~ | ✓ | https://github.com/pytorch/pytorch/issues/58734 -# uint64 | ✓ | ✓ | ~ | ✓ | https://github.com/pytorch/pytorch/issues/58734 - -# float128 | ✓ | ✗ | ✗ | ✗ | platform-dependent -# complex256 | ✓ | ✗ | ✗ | ✗ | platform-dependent - -# qint8 | ✗ | ✗ | ✓ | ✓ | -# qint16 | ✗ | ✗ | ✗ | ✓ | -# qint32 | ✗ | ✗ | ✓ | ✓ | - -# quint8 | ✗ | ✗ | ✓ | ✓ | -# quint16 | ✗ | ✗ | ✗ | ✓ | - -# bfloat16 | ✗* | ✓ | ✓ | ✓ | * available with ml_dtypes dependency - -# unicode_ | ✓ | ✗ | ✗ | ✗ | np.str_ -# bytes_ | ✓ | ✗ | ✗ | ✓ | np.bytes_, tf.string -# object | ✓ | ✗ | ✗ | ✗ | -# void | ✓ | ✗ | ✗ | ✗ | - - -# dict for reverse search native dtype -> iop dtype -_NATIVE_TO_IOP: dict[object, dtype] = {} - - -def _initialize_dtype_registry(backend: Backend) -> None: - """Initialize dtype objects and the reverse native-to-iop map for ``backend``.""" - global _NATIVE_TO_IOP # noqa: PLW0603 - module = sys.modules[__name__] - _NATIVE_TO_IOP = {} - - for name in (_SUPPORTED_DTYPES + _OPTIONAL_DTYPES): - _register_dtype(module, backend, name) - - -def _register_dtype(module: object, backend: Backend, name: str) -> None: - """Create and publish a dtype (if available).""" - native_dtype = getattr(backend, name, None) - if native_dtype is None: - return - - iop_dtype = dtype(name) - setattr(module, name, iop_dtype) - _NATIVE_TO_IOP[native_dtype] = iop_dtype diff --git a/decent_array/types/__init__.py b/decent_array/types/__init__.py new file mode 100644 index 0000000..2e77581 --- /dev/null +++ b/decent_array/types/__init__.py @@ -0,0 +1,51 @@ +from typing import TYPE_CHECKING + +from decent_array.types._dtypes import _ALL_DTYPES, dtype, dtypes +from decent_array.types._types import ( + ArrayKey as ArrayKey, +) +from decent_array.types._types import ( + ArrayLike as ArrayLike, +) +from decent_array.types._types import ( + DTypes as DTypes, +) +from decent_array.types._types import ( + SupportedArrayTypes as SupportedArrayTypes, +) +from decent_array.types._types import ( + SupportedDevices as SupportedDevices, +) +from decent_array.types._types import ( + SupportedFrameworks as SupportedFrameworks, +) + +_static_exports = [ # noqa: RUF067 + "ArrayLike", + "SupportedArrayTypes", + "ArrayKey", + "SupportedFrameworks", + "SupportedDevices", + "DTypes", + "dtype", + "dtypes", +] + +_all_dtypes = list(_ALL_DTYPES.keys()) # noqa: RUF067 +_available_dtypes = [name for name, dt in _ALL_DTYPES.items() if dt.available] # noqa: RUF067 + +__all_docs__ = _static_exports + _all_dtypes + +__all__ = [ + "ArrayKey", + "ArrayLike", + "DTypes", + "SupportedArrayTypes", + "SupportedDevices", + "SupportedFrameworks", + "dtype", + "dtypes", +] + +if not TYPE_CHECKING: # noqa: RUF067 + __all__ += _available_dtypes # noqa: PLE0605 diff --git a/decent_array/types/_dtypes.py b/decent_array/types/_dtypes.py new file mode 100644 index 0000000..0a8126c --- /dev/null +++ b/decent_array/types/_dtypes.py @@ -0,0 +1,141 @@ +"""Data type definitions.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +from decent_array.interoperability._backend_manager import register_backend_listener + +if TYPE_CHECKING: + from decent_array.interoperability._abstracts import Backend + + +_BACKEND_INSTANCE: Backend | None = None +_error = RuntimeError("No backend active: call 'set_backend' with a supported framework to activate one.") + + +def _update_backend(backend: Backend | None) -> None: + global _BACKEND_INSTANCE # noqa: PLW0603 + _BACKEND_INSTANCE = backend + + +register_backend_listener(_update_backend) + + +class dtype: # noqa: N801 + """Base class for dtypes.""" + + def __init__(self, name: str): + if _BACKEND_INSTANCE is None: + raise _error + + # name doesn't map to any dtype + if name not in _ALL_DTYPES: + raise ValueError(f"dtype {name} is not supported.") + + # this is None if the dtype is not supported + backend_dtype = getattr(_BACKEND_INSTANCE, name, None) + + self._name = name + self._backend_dtype: Any = backend_dtype + self._available = backend_dtype is not None + + @property + def name(self) -> str: + """Name of the dtype.""" + return self._name + + @property + def available(self) -> bool: + """Availability of the dtype (dependent on backend, device, backend settings, OS).""" + return self._available + + def __str__(self) -> str: + """Name of the dtype.""" + return self.name + + def __eq__(self, other: object) -> bool: + """Check equivalence by ``name`` attributes.""" + if not isinstance(other, dtype): + return NotImplemented + return self.name == other.name + + def __hash__(self) -> int: + """Hash of the dtype.""" + return hash(self.name) + + +# TODO create global error message for this when an unavailable dtype is attampted in astype and other funcs +# raise ValueError(f"dtype {name} is not supported.") # TODO add device and backend name to contextualize + +# TODO do we want to have an optional dependence on ml_dtypes to support bfloat16 for numpy? + + +_BOOL_DTYPES = { + "bool_": dtype("bool"), +} + +_SIGNED_INT_DTYPES = { + "int8": dtype("int8"), + "int16": dtype("int16"), + "int32": dtype("int32"), + "int64": dtype("int64"), +} + +_UNSIGNED_INT_DTYPES = { + "uint8": dtype("uint8"), + "uint16": dtype("uint16"), + "uint32": dtype("uint32"), + "uint64": dtype("uint64"), +} + +_REAL_FLOATING_DTYPES = { + "float16": dtype("float16"), + "bfloat16": dtype("bfloat16"), + "float32": dtype("float32"), + "float64": dtype("float64"), + "float128": dtype("float128"), +} + +_COMPLEX_FLOATING_DTYPES = { + "complex64": dtype("complex64"), + "complex128": dtype("complex128"), + "complex256": dtype("complex256"), +} + +_QUANTIZED_SIGNED_INT_DTYPES = { + "qint8": dtype("qint8"), + "qint16": dtype("qint16"), + "qint32": dtype("qint32"), +} + +_QUANTIZED_UNSIGNED_INT_DTYPES = { + "quint8": dtype("quint8"), + "quint16": dtype("quint16"), +} + +_MISCELLANEOUS_DTYPES = { + "unicode_": dtype("unicode"), + "bytes_": dtype("bytes"), + "object": dtype("object"), + "void": dtype("void"), +} + +_INTEGRAL_DTYPES = (_SIGNED_INT_DTYPES | + _UNSIGNED_INT_DTYPES | + _QUANTIZED_SIGNED_INT_DTYPES | + _QUANTIZED_UNSIGNED_INT_DTYPES + ) +_NUMERIC_DTYPES = _INTEGRAL_DTYPES | _REAL_FLOATING_DTYPES | _COMPLEX_FLOATING_DTYPES +_ALL_DTYPES = _BOOL_DTYPES | _NUMERIC_DTYPES | _MISCELLANEOUS_DTYPES + + +_BACKEND_DTYPE_TO_DTYPE: dict[Any, dtype] = {} +for dt in _ALL_DTYPES.values(): + if dt.available: + _BACKEND_DTYPE_TO_DTYPE[dt._backend_dtype] = dt # noqa: SLF001 + + +def dtypes() -> dict[str, dtype]: + """Return a dictionary of available dtypes.""" + return {name: dt for name, dt in _ALL_DTYPES.items() if dt.available} diff --git a/decent_array/types/_types.py b/decent_array/types/_types.py new file mode 100644 index 0000000..74f00e0 --- /dev/null +++ b/decent_array/types/_types.py @@ -0,0 +1,91 @@ +"""Type definitions.""" + +from __future__ import annotations + +from enum import Enum +from typing import TYPE_CHECKING, SupportsIndex, TypeAlias, Union + +from decent_array.interoperability._backend_manager import register_backend_listener + +if TYPE_CHECKING: + import jax + import numpy + import tensorflow as tf + import torch + + from decent_array._array import Array + from decent_array.interoperability._abstracts import Backend + + +_BACKEND_INSTANCE: Backend | None = None +_error = RuntimeError("No backend active: call 'set_backend' with a supported framework to activate one.") + + +def _update_backend(backend: Backend | None) -> None: + global _BACKEND_INSTANCE # noqa: PLW0603 + _BACKEND_INSTANCE = backend + + +register_backend_listener(_update_backend) + + +ArrayLike: TypeAlias = Union["numpy.ndarray", "torch.Tensor", "tf.Tensor", "jax.Array"] # noqa: UP040 +""" +Type alias for array-like types supported in decent-array, including NumPy arrays, +PyTorch tensors, TensorFlow tensors, and JAX arrays. +""" + +SupportedArrayTypes: TypeAlias = bool | int | float | complex | ArrayLike # noqa: UP040 +""" +Type alias for supported types for optimization variables in decent-array, +including array-like types and scalars. +""" + +ArrayKey: TypeAlias = ( # noqa: UP040 + "int | SupportsIndex | slice | Array | tuple[int | SupportsIndex | slice | Array | None, ...] | None" +) +""" +Type alias for valid keys used to index into supported array types. +Includes single indices, tuples of indices, slices, and tuples of slices. +""" + + +# Its important that the enum values correspond to the folder names of the backends, +# since those are used for dynamic imports in _backend_manager.py +class SupportedFrameworks(Enum): + """Enum for supported frameworks in decent-array.""" + + NUMPY = "numpy" + PYTORCH = "pytorch" + TENSORFLOW = "tensorflow" + JAX = "jax" + + +class SupportedDevices(Enum): + """Enum for supported devices in decent-array.""" + + CPU = "cpu" + GPU = "gpu" + MPS = "mps" + + +class DTypes(Enum): + """Enum for supported dtypes in decent-array.""" + + BOOL = "bool" + UINT8 = "uint8" + UINT16 = "uint16" + UINT32 = "uint32" + UINT64 = "uint64" + INT8 = "int8" + INT16 = "int16" + INT32 = "int32" + INT64 = "int64" + FLOAT16 = "float16" + FLOAT32 = "float32" + FLOAT64 = "float64" + COMPLEX64 = "complex64" + COMPLEX128 = "complex128" + + +_STRING_TO_DTYPE = {dt.value: dt for dt in DTypes} From b987a8218f955b783ce6e4ae0bec111bd6ca4d07 Mon Sep 17 00:00:00 2001 From: nicola-bastianello Date: Mon, 15 Jun 2026 17:46:11 +0200 Subject: [PATCH 04/23] ref: change enums names --- decent_array/_array.py | 6 +-- .../interoperability/_abstracts/backend.py | 20 ++++---- .../interoperability/_backend_manager.py | 38 +++++++-------- .../interoperability/_iop/manipulations.py | 4 +- decent_array/interoperability/_iop/rng.py | 4 +- decent_array/interoperability/_iop/utils.py | 10 ++-- .../interoperability/_jax/jax_backend.py | 20 ++++---- .../interoperability/_numpy/numpy_backend.py | 16 +++---- .../_pytorch/pytorch_backend.py | 24 +++++----- .../_tensorflow/tensorflow_backend.py | 18 ++++---- decent_array/types/__init__.py | 18 ++++---- decent_array/types/_types.py | 6 +-- tests/conftest.py | 46 +++++++++---------- tests/test_backend_manager.py | 40 ++++++++-------- 14 files changed, 135 insertions(+), 135 deletions(-) diff --git a/decent_array/_array.py b/decent_array/_array.py index 4a0224a..83a36f4 100644 --- a/decent_array/_array.py +++ b/decent_array/_array.py @@ -35,7 +35,7 @@ from numpy.typing import NDArray from decent_array.interoperability._abstracts import Backend - from decent_array.types import ArrayKey, DTypes, SupportedArrayTypes, SupportedDevices + from decent_array.types import ArrayKey, DTypes, ArrayTypes, Devices _BACKEND_INSTANCE: Backend | None = None @@ -60,7 +60,7 @@ class Array: # noqa: PLR0904 __slots__ = ("_backend", "value") - def __init__(self, value: SupportedArrayTypes) -> None: + def __init__(self, value: ArrayTypes) -> None: """ Wrap ``value`` in an :class:`Array`. @@ -445,7 +445,7 @@ def all(self) -> bool: return self._backend.all(self) @property - def device(self) -> SupportedDevices: + def device(self) -> Devices: """Return the device of the array.""" return self._backend.device_of(self) diff --git a/decent_array/interoperability/_abstracts/backend.py b/decent_array/interoperability/_abstracts/backend.py index edf7e6e..21f4e54 100644 --- a/decent_array/interoperability/_abstracts/backend.py +++ b/decent_array/interoperability/_abstracts/backend.py @@ -18,25 +18,25 @@ from collections.abc import Sequence from typing import TYPE_CHECKING, Any -from decent_array.types import SupportedDevices +from decent_array.types import Devices if TYPE_CHECKING: from numpy.typing import NDArray from decent_array import Array - from decent_array.types import ArrayKey, DTypes, SupportedArrayTypes + from decent_array.types import ArrayKey, DTypes, ArrayTypes class Backend(ABC): # noqa: PLR0904 """ Abstract base class for a backend. - Concrete backends are bound to a single :class:`SupportedDevices` at construction + Concrete backends are bound to a single :class:`Devices` at construction time; that device is the default for all new arrays produced by this backend. """ - def __init__(self, device: SupportedDevices = SupportedDevices.CPU) -> None: - self.device: SupportedDevices = device + def __init__(self, device: Devices = Devices.CPU) -> None: + self.device: Devices = device # Array creation ------------------------------------------------------ @@ -61,12 +61,12 @@ def eye(self, n: int) -> Array: """Create an ``n x n`` identity matrix.""" @abstractmethod - def device_to_native(self, device: SupportedDevices) -> Any: # noqa: ANN401 - """Convert :class:`SupportedDevices` to the backend's native device representation.""" + def device_to_native(self, device: Devices) -> Any: # noqa: ANN401 + """Convert :class:`Devices` to the backend's native device representation.""" @abstractmethod - def device_of(self, x: Array) -> SupportedDevices: - """Return the :class:`SupportedDevices` of the given array.""" + def device_of(self, x: Array) -> Devices: + """Return the :class:`Devices` of the given array.""" # Array manipulation -------------------------------------------------- @@ -75,7 +75,7 @@ def copy(self, x: Array) -> Array: """Return a copy of ``x``.""" @abstractmethod - def to_numpy(self, x: SupportedArrayTypes | Array) -> NDArray[Any]: + def to_numpy(self, x: ArrayTypes | Array) -> NDArray[Any]: """Convert ``x`` to a NumPy array on the CPU.""" @abstractmethod diff --git a/decent_array/interoperability/_backend_manager.py b/decent_array/interoperability/_backend_manager.py index a359a67..2c2621c 100644 --- a/decent_array/interoperability/_backend_manager.py +++ b/decent_array/interoperability/_backend_manager.py @@ -4,13 +4,13 @@ from collections.abc import Callable from contextvars import ContextVar -from decent_array.types import SupportedDevices, SupportedFrameworks +from decent_array.types import Devices, Frameworks from ._abstracts import Backend -_BACKEND_REGISTRY: dict[SupportedFrameworks, type[Backend]] = {} -_BACKEND_INSTANCES: dict[SupportedFrameworks, Backend] = {} -_ACTIVE_BACKEND: ContextVar[SupportedFrameworks | None] = ContextVar( +_BACKEND_REGISTRY: dict[Frameworks, type[Backend]] = {} +_BACKEND_INSTANCES: dict[Frameworks, Backend] = {} +_ACTIVE_BACKEND: ContextVar[Frameworks | None] = ContextVar( "decent_array.interoperability.active_backend", default=None ) _BACKEND_LISTENERS: list[Callable[[Backend | None], None]] = [] @@ -18,8 +18,8 @@ def set_backend( - backend: SupportedFrameworks | str, - device: SupportedDevices | str = SupportedDevices.CPU, + backend: Frameworks | str, + device: Devices | str = Devices.CPU, ) -> None: """ Set the active backend (and target device) for the current execution context. @@ -33,9 +33,9 @@ def set_backend( Backend modules are auto-imported on demand. Args: - backend: A :class:`~decent_array.types.SupportedFrameworks` value, its canonical string (e.g. + backend: A :class:`~decent_array.types.Frameworks` value, its canonical string (e.g. ``"numpy"``, ``"pytorch"``). - device: Target accelerator. Accepts a :class:`~decent_array.types.SupportedDevices` value or its + device: Target accelerator. Accepts a :class:`~decent_array.types.Devices` value or its string equivalent (``"cpu"``, ``"gpu"``, ``"mps"``). Defaults to CPU. The backend's array-creation methods produce arrays on this device by default. @@ -46,7 +46,7 @@ def set_backend( """ # noqa: DOC502 requested = _normalize(backend) - requested_device = device if isinstance(device, SupportedDevices) else SupportedDevices(device) + requested_device = device if isinstance(device, Devices) else Devices(device) current = _ACTIVE_BACKEND.get() if current is not None and current != requested: @@ -87,11 +87,11 @@ def register_backend_listener(listener: Callable[[Backend | None], None]) -> Non def register_backend( - backend: SupportedFrameworks, + backend: Frameworks, cls: type[Backend], ) -> None: """ - Register a backend class under a :class:`SupportedFrameworks` value. + Register a backend class under a :class:`Frameworks` value. Called once per backend module *after* the class definition. Backends are instantiated lazily on first use. Re-registering replaces the @@ -127,7 +127,7 @@ def reset_backends() -> None: listener(None) -def default_device() -> SupportedDevices: +def default_device() -> Devices: """ Return the default device for the active backend. @@ -143,24 +143,24 @@ def default_device() -> SupportedDevices: return backend.device -def _normalize(backend: SupportedFrameworks | str) -> SupportedFrameworks: +def _normalize(backend: Frameworks | str) -> Frameworks: """ - Convert a backend identifier to its canonical :class:`SupportedFrameworks` value. + Convert a backend identifier to its canonical :class:`Frameworks` value. Raises: KeyError: If the input is not a valid backend identifier. """ - if isinstance(backend, SupportedFrameworks): + if isinstance(backend, Frameworks): return backend try: - return SupportedFrameworks(backend) + return Frameworks(backend) except ValueError as exc: - valid = ", ".join(f.value for f in SupportedFrameworks) + valid = ", ".join(f.value for f in Frameworks) raise KeyError(f"Unknown backend '{backend}'. Valid backends: {valid}.") from exc -def _instantiate(backend: SupportedFrameworks, device: SupportedDevices) -> Backend: +def _instantiate(backend: Frameworks, device: Devices) -> Backend: """ Get or create a backend instance for the given backend and device. @@ -185,7 +185,7 @@ def _instantiate(backend: SupportedFrameworks, device: SupportedDevices) -> Back return instance -def _auto_import(backend: SupportedFrameworks) -> None: +def _auto_import(backend: Frameworks) -> None: """ Import the backend's package so its registration side-effect runs. diff --git a/decent_array/interoperability/_iop/manipulations.py b/decent_array/interoperability/_iop/manipulations.py index c2450e2..1abcd94 100644 --- a/decent_array/interoperability/_iop/manipulations.py +++ b/decent_array/interoperability/_iop/manipulations.py @@ -23,7 +23,7 @@ from decent_array import Array from decent_array.interoperability._abstracts import Backend - from decent_array.types import DTypes, SupportedArrayTypes + from decent_array.types import DTypes, ArrayTypes _BACKEND_INSTANCE: Backend | None = None _error = RuntimeError("No backend active: call 'set_backend' with a supported framework to activate one.") @@ -44,7 +44,7 @@ def copy(x: Array) -> Array: return _BACKEND_INSTANCE.copy(x) -def to_numpy(x: SupportedArrayTypes | Array) -> NDArray[Any]: +def to_numpy(x: ArrayTypes | Array) -> NDArray[Any]: """Convert ``x`` to a NumPy array on CPU.""" if _BACKEND_INSTANCE is None: raise _error diff --git a/decent_array/interoperability/_iop/rng.py b/decent_array/interoperability/_iop/rng.py index 6968ddb..3df90ca 100644 --- a/decent_array/interoperability/_iop/rng.py +++ b/decent_array/interoperability/_iop/rng.py @@ -21,7 +21,7 @@ from typing import TYPE_CHECKING, Any, cast from decent_array.interoperability._backend_manager import _instantiate, register_backend_listener -from decent_array.types import SupportedDevices, SupportedFrameworks +from decent_array.types import Devices, Frameworks if TYPE_CHECKING: import numpy @@ -118,7 +118,7 @@ def set_rng_state(self, state: dict[str, Any]) -> None: active.set_rng_state(state) def numpy_backend(self) -> NumpyBackend: - return cast("NumpyBackend", _instantiate(SupportedFrameworks.NUMPY, SupportedDevices.CPU)) + return cast("NumpyBackend", _instantiate(Frameworks.NUMPY, Devices.CPU)) _COORDINATOR = _RngCoordinator() diff --git a/decent_array/interoperability/_iop/utils.py b/decent_array/interoperability/_iop/utils.py index e1b737a..ab69d22 100644 --- a/decent_array/interoperability/_iop/utils.py +++ b/decent_array/interoperability/_iop/utils.py @@ -20,7 +20,7 @@ if TYPE_CHECKING: from decent_array import Array from decent_array.interoperability._abstracts import Backend - from decent_array.types import ArrayKey, SupportedDevices + from decent_array.types import ArrayKey, Devices _BACKEND_INSTANCE: Backend | None = None _error = RuntimeError("No backend active: call 'set_backend' with a supported framework to activate one.") @@ -34,15 +34,15 @@ def _update_backend(backend: Backend | None) -> None: register_backend_listener(_update_backend) -def device_to_native(device: SupportedDevices) -> Any: # noqa: ANN401 - """Convert :class:`~decent_array.types.SupportedDevices` to the active backend's native device.""" +def device_to_native(device: Devices) -> Any: # noqa: ANN401 + """Convert :class:`~decent_array.types.Devices` to the active backend's native device.""" if _BACKEND_INSTANCE is None: raise _error return _BACKEND_INSTANCE.device_to_native(device) -def device_of(x: Array) -> SupportedDevices: - """Return the :class:`~decent_array.types.SupportedDevices` of ``x``.""" +def device_of(x: Array) -> Devices: + """Return the :class:`~decent_array.types.Devices` of ``x``.""" if _BACKEND_INSTANCE is None: raise _error return _BACKEND_INSTANCE.device_of(x) diff --git a/decent_array/interoperability/_jax/jax_backend.py b/decent_array/interoperability/_jax/jax_backend.py index 677ff41..312007b 100644 --- a/decent_array/interoperability/_jax/jax_backend.py +++ b/decent_array/interoperability/_jax/jax_backend.py @@ -22,7 +22,7 @@ from decent_array import Array from decent_array.interoperability._abstracts import Backend from decent_array.interoperability._backend_manager import register_backend -from decent_array.types import ArrayKey, DTypes, SupportedArrayTypes, SupportedDevices, SupportedFrameworks +from decent_array.types import ArrayKey, DTypes, ArrayTypes, Devices, Frameworks def _unwrap(x: Any) -> Any: # noqa: ANN401 @@ -50,7 +50,7 @@ def _unwrap(x: Any) -> Any: # noqa: ANN401 class JaxBackend(Backend): # noqa: PLR0904 """JAX implementation of :class:`Backend`.""" - def __init__(self, device: SupportedDevices = SupportedDevices.CPU) -> None: + def __init__(self, device: Devices = Devices.CPU) -> None: super().__init__(device) self._native_device: jax.Device = self.device_to_native(device) self._key: jax.Array = jax.random.key(time_ns()) @@ -72,19 +72,19 @@ def ones_like(self, x: Array) -> Array: def eye(self, n: int) -> Array: return Array(jnp.eye(n, device=self._native_device)) - def device_to_native(self, device: SupportedDevices) -> jax.Device: - if device == SupportedDevices.CPU: + def device_to_native(self, device: Devices) -> jax.Device: + if device == Devices.CPU: return jax.devices("cpu")[0] - if device == SupportedDevices.GPU: + if device == Devices.GPU: return jax.devices("gpu")[0] raise ValueError(f"Unsupported device for JAX: {device}") - def device_of(self, x: Array) -> SupportedDevices: + def device_of(self, x: Array) -> Devices: platform = x.value.device.platform if platform == "gpu": - return SupportedDevices.GPU + return Devices.GPU if platform == "cpu": - return SupportedDevices.CPU + return Devices.CPU raise TypeError(f"Unsupported JAX platform: {platform}") # Array manipulation @@ -92,7 +92,7 @@ def device_of(self, x: Array) -> SupportedDevices: def copy(self, x: Array) -> Array: return Array(jnp.array(x.value, copy=True)) - def to_numpy(self, x: SupportedArrayTypes | Array) -> NDArray[Any]: + def to_numpy(self, x: ArrayTypes | Array) -> NDArray[Any]: return np.array(x.value if type(x) is Array else x) def from_numpy(self, x: NDArray[Any]) -> Array: @@ -459,4 +459,4 @@ def _x64_enabled() -> bool: return bool(_JAX_CONFIG_READ("jax_enable_x64")) -register_backend(SupportedFrameworks.JAX, JaxBackend) +register_backend(Frameworks.JAX, JaxBackend) diff --git a/decent_array/interoperability/_numpy/numpy_backend.py b/decent_array/interoperability/_numpy/numpy_backend.py index 849fd7c..0d9384c 100644 --- a/decent_array/interoperability/_numpy/numpy_backend.py +++ b/decent_array/interoperability/_numpy/numpy_backend.py @@ -17,7 +17,7 @@ from decent_array import Array from decent_array.interoperability._abstracts import Backend from decent_array.interoperability._backend_manager import register_backend -from decent_array.types import ArrayKey, DTypes, SupportedArrayTypes, SupportedDevices, SupportedFrameworks +from decent_array.types import ArrayKey, DTypes, ArrayTypes, Devices, Frameworks def _unwrap(x: Any) -> Any: # noqa: ANN401 @@ -51,8 +51,8 @@ def _unwrap(x: Any) -> Any: # noqa: ANN401 class NumpyBackend(Backend): # noqa: PLR0904 """NumPy implementation of :class:`Backend`.""" - def __init__(self, device: SupportedDevices = SupportedDevices.CPU) -> None: - if device != SupportedDevices.CPU: + def __init__(self, device: Devices = Devices.CPU) -> None: + if device != Devices.CPU: raise ValueError(f"NumPy backend only supports CPU, got '{device.value}'.") super().__init__(device) self._rng: np.random.Generator = np.random.default_rng() @@ -74,12 +74,12 @@ def ones_like(self, x: Array) -> Array: def eye(self, n: int) -> Array: return Array(np.eye(n)) - def device_to_native(self, device: SupportedDevices) -> Any: # noqa: ANN401 + def device_to_native(self, device: Devices) -> Any: # noqa: ANN401 # NumPy has no explicit device management; surface the request unchanged. return device - def device_of(self, x: Array) -> SupportedDevices: # noqa: ARG002 - return SupportedDevices.CPU + def device_of(self, x: Array) -> Devices: # noqa: ARG002 + return Devices.CPU # Array manipulation @@ -89,7 +89,7 @@ def copy(self, x: Array) -> Array: return Array(np.copy(v)) return Array(deepcopy(v)) - def to_numpy(self, x: SupportedArrayTypes | Array) -> NDArray[Any]: + def to_numpy(self, x: ArrayTypes | Array) -> NDArray[Any]: """Return the value of an :class:`Array` as a NumPy array.""" v = x.value if type(x) is Array else x if isinstance(v, np.ndarray): @@ -460,4 +460,4 @@ def void(self) -> Any: # noqa: ANN401 return np.dtype(np.void) -register_backend(SupportedFrameworks.NUMPY, NumpyBackend) +register_backend(Frameworks.NUMPY, NumpyBackend) diff --git a/decent_array/interoperability/_pytorch/pytorch_backend.py b/decent_array/interoperability/_pytorch/pytorch_backend.py index f34d241..9112fcb 100644 --- a/decent_array/interoperability/_pytorch/pytorch_backend.py +++ b/decent_array/interoperability/_pytorch/pytorch_backend.py @@ -17,7 +17,7 @@ from decent_array import Array from decent_array.interoperability._abstracts import Backend from decent_array.interoperability._backend_manager import register_backend -from decent_array.types import ArrayKey, DTypes, SupportedArrayTypes, SupportedDevices, SupportedFrameworks +from decent_array.types import ArrayKey, DTypes, ArrayTypes, Devices, Frameworks def _unwrap(x: Any) -> Any: # noqa: ANN401 @@ -45,7 +45,7 @@ def _unwrap(x: Any) -> Any: # noqa: ANN401 class PyTorchBackend(Backend): # noqa: PLR0904 """PyTorch implementation of :class:`Backend`.""" - def __init__(self, device: SupportedDevices = SupportedDevices.CPU) -> None: + def __init__(self, device: Devices = Devices.CPU) -> None: super().__init__(device) self._native_device: str = self.device_to_native(device) self._generator: torch.Generator = torch.Generator(device=self._native_device) @@ -67,23 +67,23 @@ def ones_like(self, x: Array) -> Array: def eye(self, n: int) -> Array: return Array(torch.eye(n, device=self._native_device)) - def device_to_native(self, device: SupportedDevices) -> str: - if device == SupportedDevices.CPU: + def device_to_native(self, device: Devices) -> str: + if device == Devices.CPU: return "cpu" - if device == SupportedDevices.GPU: + if device == Devices.GPU: return "cuda" - if device == SupportedDevices.MPS: + if device == Devices.MPS: return "mps" raise ValueError(f"Unsupported device: {device}") - def device_of(self, x: Array) -> SupportedDevices: + def device_of(self, x: Array) -> Devices: kind = x.value.device.type if kind == "cpu": - return SupportedDevices.CPU + return Devices.CPU if kind == "cuda": - return SupportedDevices.GPU + return Devices.GPU if kind == "mps": - return SupportedDevices.MPS + return Devices.MPS raise TypeError(f"Unsupported PyTorch device type: {kind}") # Array manipulation @@ -91,7 +91,7 @@ def device_of(self, x: Array) -> SupportedDevices: def copy(self, x: Array) -> Array: return Array(x.value.detach().clone()) - def to_numpy(self, x: SupportedArrayTypes | Array) -> NDArray[Any]: + def to_numpy(self, x: ArrayTypes | Array) -> NDArray[Any]: """Return the value of an :class:`Array` as a NumPy array.""" v = x.value if type(x) is Array else x if isinstance(v, torch.Tensor): @@ -485,4 +485,4 @@ def bfloat16(self) -> torch.dtype: return torch.bfloat16 -register_backend(SupportedFrameworks.PYTORCH, PyTorchBackend) +register_backend(Frameworks.PYTORCH, PyTorchBackend) diff --git a/decent_array/interoperability/_tensorflow/tensorflow_backend.py b/decent_array/interoperability/_tensorflow/tensorflow_backend.py index d39dac3..abf3654 100644 --- a/decent_array/interoperability/_tensorflow/tensorflow_backend.py +++ b/decent_array/interoperability/_tensorflow/tensorflow_backend.py @@ -20,7 +20,7 @@ from decent_array import Array from decent_array.interoperability._abstracts import Backend from decent_array.interoperability._backend_manager import register_backend -from decent_array.types import ArrayKey, DTypes, SupportedArrayTypes, SupportedDevices, SupportedFrameworks +from decent_array.types import ArrayKey, DTypes, ArrayTypes, Devices, Frameworks def _unwrap(x: Any) -> Any: # noqa: ANN401 @@ -49,7 +49,7 @@ def _unwrap(x: Any) -> Any: # noqa: ANN401 class TensorflowBackend(Backend): # noqa: PLR0904 """TensorFlow implementation of :class:`Backend`.""" - def __init__(self, device: SupportedDevices = SupportedDevices.CPU) -> None: + def __init__(self, device: Devices = Devices.CPU) -> None: super().__init__(device) self._native_device: str = self.device_to_native(device) self._generator: tf.random.Generator = tf.random.Generator.from_non_deterministic_state(alg="philox") @@ -74,23 +74,23 @@ def eye(self, n: int) -> Array: with tf.device(self._native_device): return Array(tf.eye(n)) - def device_to_native(self, device: SupportedDevices) -> str: - if device in {SupportedDevices.CPU, SupportedDevices.GPU}: + def device_to_native(self, device: Devices) -> str: + if device in {Devices.CPU, Devices.GPU}: return f"/{device.value}:0" raise ValueError(f"Unsupported device for TensorFlow: {device}") - def device_of(self, x: Array) -> SupportedDevices: + def device_of(self, x: Array) -> Devices: device_str = x.value.device.lower() if "gpu" in device_str or "cuda" in device_str: - return SupportedDevices.GPU - return SupportedDevices.CPU + return Devices.GPU + return Devices.CPU # Array manipulation def copy(self, x: Array) -> Array: return Array(tf.identity(x.value)) - def to_numpy(self, x: SupportedArrayTypes | Array) -> NDArray[Any]: + def to_numpy(self, x: ArrayTypes | Array) -> NDArray[Any]: """Return the value of an :class:`Array` as a NumPy array.""" v = x.value if type(x) is Array else x if isinstance(v, tf.Tensor): @@ -517,4 +517,4 @@ def bfloat16(self) -> tf.dtypes.DType: return tf.bfloat16 -register_backend(SupportedFrameworks.TENSORFLOW, TensorflowBackend) +register_backend(Frameworks.TENSORFLOW, TensorflowBackend) diff --git a/decent_array/types/__init__.py b/decent_array/types/__init__.py index 2e77581..e69df9e 100644 --- a/decent_array/types/__init__.py +++ b/decent_array/types/__init__.py @@ -11,21 +11,21 @@ DTypes as DTypes, ) from decent_array.types._types import ( - SupportedArrayTypes as SupportedArrayTypes, + ArrayTypes as ArrayTypes, ) from decent_array.types._types import ( - SupportedDevices as SupportedDevices, + Devices as Devices, ) from decent_array.types._types import ( - SupportedFrameworks as SupportedFrameworks, + Frameworks as Frameworks, ) _static_exports = [ # noqa: RUF067 "ArrayLike", - "SupportedArrayTypes", + "ArrayTypes", "ArrayKey", - "SupportedFrameworks", - "SupportedDevices", + "Frameworks", + "Devices", "DTypes", "dtype", "dtypes", @@ -40,9 +40,9 @@ "ArrayKey", "ArrayLike", "DTypes", - "SupportedArrayTypes", - "SupportedDevices", - "SupportedFrameworks", + "ArrayTypes", + "Devices", + "Frameworks", "dtype", "dtypes", ] diff --git a/decent_array/types/_types.py b/decent_array/types/_types.py index 74f00e0..963cabd 100644 --- a/decent_array/types/_types.py +++ b/decent_array/types/_types.py @@ -35,7 +35,7 @@ def _update_backend(backend: Backend | None) -> None: PyTorch tensors, TensorFlow tensors, and JAX arrays. """ -SupportedArrayTypes: TypeAlias = bool | int | float | complex | ArrayLike # noqa: UP040 +ArrayTypes: TypeAlias = bool | int | float | complex | ArrayLike # noqa: UP040 """ Type alias for supported types for optimization variables in decent-array, including array-like types and scalars. @@ -52,7 +52,7 @@ def _update_backend(backend: Backend | None) -> None: # Its important that the enum values correspond to the folder names of the backends, # since those are used for dynamic imports in _backend_manager.py -class SupportedFrameworks(Enum): +class Frameworks(Enum): """Enum for supported frameworks in decent-array.""" NUMPY = "numpy" @@ -61,7 +61,7 @@ class SupportedFrameworks(Enum): JAX = "jax" -class SupportedDevices(Enum): +class Devices(Enum): """Enum for supported devices in decent-array.""" CPU = "cpu" diff --git a/tests/conftest.py b/tests/conftest.py index 1ca38bf..8e83362 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,7 +1,7 @@ """Shared fixtures: parametrize tests across every (framework, device) combination. Each test using the ``backend`` fixture runs once per (framework, device) pair from -:class:`SupportedFrameworks` x :class:`SupportedDevices`. Combinations whose backend +:class:`Frameworks` x :class:`Devices`. Combinations whose backend package is missing or whose device is not present on the current host are marked ``skip`` so the test report stays interpretable on machines with partial accelerator support. @@ -15,50 +15,50 @@ import pytest from decent_array.interoperability._backend_manager import reset_backends -from decent_array.types import SupportedDevices, SupportedFrameworks +from decent_array.types import Devices, Frameworks if TYPE_CHECKING: from _pytest.fixtures import FixtureRequest -def _framework_importable(framework: SupportedFrameworks) -> bool: +def _framework_importable(framework: Frameworks) -> bool: try: - if framework == SupportedFrameworks.NUMPY: + if framework == Frameworks.NUMPY: import numpy # noqa: F401, PLC0415 - elif framework == SupportedFrameworks.PYTORCH: + elif framework == Frameworks.PYTORCH: import torch # noqa: F401, PLC0415 - elif framework == SupportedFrameworks.JAX: + elif framework == Frameworks.JAX: import jax # noqa: F401, PLC0415 - elif framework == SupportedFrameworks.TENSORFLOW: + elif framework == Frameworks.TENSORFLOW: import tensorflow # noqa: F401, PLC0415 except ImportError: return False return True -def _device_available(framework: SupportedFrameworks, device: SupportedDevices) -> bool: +def _device_available(framework: Frameworks, device: Devices) -> bool: """Return True iff this (framework, device) pair can run on the current host.""" if not _framework_importable(framework): return False - if framework == SupportedFrameworks.NUMPY: - return device == SupportedDevices.CPU - if framework == SupportedFrameworks.PYTORCH: + if framework == Frameworks.NUMPY: + return device == Devices.CPU + if framework == Frameworks.PYTORCH: import torch # noqa: PLC0415 - if device == SupportedDevices.CPU: + if device == Devices.CPU: return True - if device == SupportedDevices.GPU: + if device == Devices.GPU: try: return bool(torch.cuda.is_available()) except Exception: return False - if device == SupportedDevices.MPS: + if device == Devices.MPS: try: return bool(torch.backends.mps.is_available()) except Exception: return False - if framework == SupportedFrameworks.JAX: - if device == SupportedDevices.MPS: + if framework == Frameworks.JAX: + if device == Devices.MPS: return False import jax # noqa: PLC0415 @@ -67,14 +67,14 @@ def _device_available(framework: SupportedFrameworks, device: SupportedDevices) except Exception: return False return True - if framework == SupportedFrameworks.TENSORFLOW: - if device == SupportedDevices.MPS: + if framework == Frameworks.TENSORFLOW: + if device == Devices.MPS: return False import tensorflow as tf # noqa: PLC0415 - if device == SupportedDevices.CPU: + if device == Devices.CPU: return True - if device == SupportedDevices.GPU: + if device == Devices.GPU: try: return len(tf.config.list_physical_devices("GPU")) > 0 except Exception: @@ -84,8 +84,8 @@ def _device_available(framework: SupportedFrameworks, device: SupportedDevices) def _backend_params() -> list[pytest.ParameterSet]: params: list[pytest.ParameterSet] = [] - for framework in SupportedFrameworks: - for device in SupportedDevices: + for framework in Frameworks: + for device in Devices: test_id = f"{framework.value}-{device.value}" if _device_available(framework, device): params.append(pytest.param((framework, device), id=test_id)) @@ -104,7 +104,7 @@ def _backend_params() -> list[pytest.ParameterSet]: @pytest.fixture(params=BACKEND_PARAMS) -def backend(request: FixtureRequest) -> Iterator[tuple[SupportedFrameworks, SupportedDevices]]: +def backend(request: FixtureRequest) -> Iterator[tuple[Frameworks, Devices]]: """Activate the (framework, device) backend for this test, then reset on teardown.""" from decent_array.interoperability import set_backend # noqa: PLC0415 diff --git a/tests/test_backend_manager.py b/tests/test_backend_manager.py index a0eaf4c..3f79e66 100644 --- a/tests/test_backend_manager.py +++ b/tests/test_backend_manager.py @@ -17,7 +17,7 @@ reset_backends, set_backend, ) -from decent_array.types import SupportedDevices, SupportedFrameworks +from decent_array.types import Devices, Frameworks if TYPE_CHECKING: from collections.abc import Iterator @@ -40,11 +40,11 @@ def _isolate_listeners_and_backends() -> Iterator[None]: def test_normalize_accepts_enum() -> None: - assert _normalize(SupportedFrameworks.NUMPY) == SupportedFrameworks.NUMPY + assert _normalize(Frameworks.NUMPY) == Frameworks.NUMPY def test_normalize_accepts_string() -> None: - assert _normalize("numpy") == SupportedFrameworks.NUMPY + assert _normalize("numpy") == Frameworks.NUMPY def test_normalize_unknown_raises() -> None: @@ -57,19 +57,19 @@ def test_normalize_unknown_raises() -> None: def test_set_backend_with_string() -> None: set_backend("numpy") - assert backend_manager._ACTIVE_BACKEND.get() == SupportedFrameworks.NUMPY + assert backend_manager._ACTIVE_BACKEND.get() == Frameworks.NUMPY def test_set_backend_with_enum() -> None: - set_backend(SupportedFrameworks.NUMPY) - assert backend_manager._ACTIVE_BACKEND.get() == SupportedFrameworks.NUMPY + set_backend(Frameworks.NUMPY) + assert backend_manager._ACTIVE_BACKEND.get() == Frameworks.NUMPY def test_set_backend_idempotent_same_backend() -> None: set_backend("numpy") # Re-activating with the same backend+device must be a no-op (no exception). set_backend("numpy") - assert backend_manager._ACTIVE_BACKEND.get() == SupportedFrameworks.NUMPY + assert backend_manager._ACTIVE_BACKEND.get() == Frameworks.NUMPY def test_set_backend_different_backend_raises() -> None: @@ -80,8 +80,8 @@ def test_set_backend_different_backend_raises() -> None: def test_set_backend_with_string_device() -> None: set_backend("numpy", "cpu") - instance = _instantiate(SupportedFrameworks.NUMPY, SupportedDevices.CPU) - assert instance.device == SupportedDevices.CPU + instance = _instantiate(Frameworks.NUMPY, Devices.CPU) + assert instance.device == Devices.CPU def test_set_backend_invalid_name_raises() -> None: @@ -97,21 +97,21 @@ class NotABackend: pass with pytest.raises(TypeError, match=r"subclass of Backend"): - register_backend(SupportedFrameworks.NUMPY, NotABackend) # type: ignore[arg-type] + register_backend(Frameworks.NUMPY, NotABackend) # type: ignore[arg-type] def test_register_backend_replaces_cached_instance() -> None: # First import registers the real backend; instantiate to populate cache. set_backend("numpy") - cached = backend_manager._BACKEND_INSTANCES.get(SupportedFrameworks.NUMPY) + cached = backend_manager._BACKEND_INSTANCES.get(Frameworks.NUMPY) assert cached is not None # Re-register the same class — cache should be cleared so next instantiate is fresh. from decent_array.interoperability._numpy.numpy_backend import NumpyBackend # noqa: PLC0415 reset_backends() - register_backend(SupportedFrameworks.NUMPY, NumpyBackend) - assert SupportedFrameworks.NUMPY not in backend_manager._BACKEND_INSTANCES + register_backend(Frameworks.NUMPY, NumpyBackend) + assert Frameworks.NUMPY not in backend_manager._BACKEND_INSTANCES # register_backend_listener --------------------------------------------- @@ -167,22 +167,22 @@ def test_reset_backends_clears_active() -> None: def test_reset_backends_clears_instance_cache() -> None: set_backend("numpy") - assert SupportedFrameworks.NUMPY in backend_manager._BACKEND_INSTANCES + assert Frameworks.NUMPY in backend_manager._BACKEND_INSTANCES reset_backends() - assert SupportedFrameworks.NUMPY not in backend_manager._BACKEND_INSTANCES + assert Frameworks.NUMPY not in backend_manager._BACKEND_INSTANCES # _instantiate ---------------------------------------------------------- def test_instantiate_caches_instance() -> None: - a = _instantiate(SupportedFrameworks.NUMPY, SupportedDevices.CPU) - b = _instantiate(SupportedFrameworks.NUMPY, SupportedDevices.CPU) + a = _instantiate(Frameworks.NUMPY, Devices.CPU) + b = _instantiate(Frameworks.NUMPY, Devices.CPU) assert a is b def test_set_backend_device_mismatch_raises() -> None: - set_backend("numpy", SupportedDevices.CPU) + set_backend("numpy", Devices.CPU) # NumPy backend rejects non-CPU devices at construction; check behavior via the # configured-mismatch path: re-set with a different device after first activation. with pytest.raises((RuntimeError, ValueError)): @@ -194,8 +194,8 @@ def test_set_backend_device_mismatch_raises() -> None: def test_default_device_returns_active_device() -> None: - set_backend("numpy", SupportedDevices.CPU) - assert default_device() == SupportedDevices.CPU + set_backend("numpy", Devices.CPU) + assert default_device() == Devices.CPU def test_default_device_raises_when_no_backend() -> None: From b963ff8f06216758a70e10a4ced1a59d68997a49 Mon Sep 17 00:00:00 2001 From: nicola-bastianello Date: Mon, 15 Jun 2026 17:51:45 +0200 Subject: [PATCH 05/23] enh: change default dtype names --- .../interoperability/_abstracts/backend.py | 16 ++-- .../interoperability/_jax/jax_backend.py | 2 +- .../interoperability/_numpy/numpy_backend.py | 8 +- .../_pytorch/pytorch_backend.py | 2 +- .../_tensorflow/tensorflow_backend.py | 2 +- decent_array/types/_dtypes.py | 83 ++++++++++--------- 6 files changed, 60 insertions(+), 53 deletions(-) diff --git a/decent_array/interoperability/_abstracts/backend.py b/decent_array/interoperability/_abstracts/backend.py index 21f4e54..3ca8e0b 100644 --- a/decent_array/interoperability/_abstracts/backend.py +++ b/decent_array/interoperability/_abstracts/backend.py @@ -398,7 +398,7 @@ def choice(self, x: Array, size: int, replace: bool = True) -> Array: @property @abstractmethod - def bool(self) -> Any: # noqa: ANN401 + def bool_(self) -> Any: # noqa: ANN401 """Bool dtype.""" @property @@ -437,9 +437,9 @@ def int32(self) -> Any: # noqa: ANN401 """Signed 32-bit integer dtype.""" @property - @abstractmethod def int64(self) -> Any: # noqa: ANN401 """Signed 64-bit integer dtype.""" + return None @property @abstractmethod @@ -452,19 +452,19 @@ def float32(self) -> Any: # noqa: ANN401 """32-bit floating-point dtype.""" @property - @abstractmethod def float64(self) -> Any: # noqa: ANN401 """64-bit floating-point dtype.""" + return None @property - @abstractmethod def complex64(self) -> Any: # noqa: ANN401 """64-bit complex dtype.""" + return None @property - @abstractmethod def complex128(self) -> Any: # noqa: ANN401 """128-bit complex dtype.""" + return None @property def float128(self) -> Any | None: # noqa: ANN401 @@ -507,17 +507,17 @@ def bfloat16(self) -> Any | None: # noqa: ANN401 return None @property - def unicode(self) -> Any | None: # noqa: ANN401 + def unicode_(self) -> Any | None: # noqa: ANN401 """Unicode string dtype (optional).""" return None @property - def bytes(self) -> Any | None: # noqa: ANN401 + def bytes_(self) -> Any | None: # noqa: ANN401 """Byte string dtype (optional).""" return None @property - def object(self) -> Any | None: # noqa: ANN401 + def object_(self) -> Any | None: # noqa: ANN401 """Python object dtype (optional).""" return None diff --git a/decent_array/interoperability/_jax/jax_backend.py b/decent_array/interoperability/_jax/jax_backend.py index 312007b..935be9e 100644 --- a/decent_array/interoperability/_jax/jax_backend.py +++ b/decent_array/interoperability/_jax/jax_backend.py @@ -384,7 +384,7 @@ def _next_key(self) -> jax.Array: # Dtypes @property - def bool(self) -> Any: # noqa: ANN401 + def bool_(self) -> Any: # noqa: ANN401 return np.dtype(jnp.bool_) @property diff --git a/decent_array/interoperability/_numpy/numpy_backend.py b/decent_array/interoperability/_numpy/numpy_backend.py index 0d9384c..062ee6f 100644 --- a/decent_array/interoperability/_numpy/numpy_backend.py +++ b/decent_array/interoperability/_numpy/numpy_backend.py @@ -378,7 +378,7 @@ def choice(self, x: Array, size: int, replace: bool = True) -> Array: # Dtypes @property - def bool(self) -> Any: # noqa: ANN401 + def bool_(self) -> Any: # noqa: ANN401 return np.dtype(np.bool_) @property @@ -444,15 +444,15 @@ def complex256(self) -> Any | None: # noqa: ANN401 return np.dtype(complex256) if complex256 is not None else None @property - def unicode(self) -> Any: # noqa: ANN401 + def unicode_(self) -> Any: # noqa: ANN401 return np.dtype(np.str_) @property - def bytes(self) -> Any: # noqa: ANN401 + def bytes_(self) -> Any: # noqa: ANN401 return np.dtype(np.bytes_) @property - def object(self) -> Any: # noqa: ANN401 + def object_(self) -> Any: # noqa: ANN401 return np.dtype(np.object_) @property diff --git a/decent_array/interoperability/_pytorch/pytorch_backend.py b/decent_array/interoperability/_pytorch/pytorch_backend.py index 9112fcb..95f0abd 100644 --- a/decent_array/interoperability/_pytorch/pytorch_backend.py +++ b/decent_array/interoperability/_pytorch/pytorch_backend.py @@ -425,7 +425,7 @@ def choice(self, x: Array, size: int, replace: bool = True) -> Array: # Dtypes @property - def bool(self) -> torch.dtype: + def bool_(self) -> torch.dtype: return torch.bool @property diff --git a/decent_array/interoperability/_tensorflow/tensorflow_backend.py b/decent_array/interoperability/_tensorflow/tensorflow_backend.py index abf3654..2bbb2a4 100644 --- a/decent_array/interoperability/_tensorflow/tensorflow_backend.py +++ b/decent_array/interoperability/_tensorflow/tensorflow_backend.py @@ -437,7 +437,7 @@ def choice(self, x: Array, size: int, replace: bool = True) -> Array: # Dtypes @property - def bool(self) -> tf.dtypes.DType: + def bool_(self) -> tf.dtypes.DType: return tf.bool @property diff --git a/decent_array/types/_dtypes.py b/decent_array/types/_dtypes.py index 0a8126c..eea5074 100644 --- a/decent_array/types/_dtypes.py +++ b/decent_array/types/_dtypes.py @@ -2,43 +2,50 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Any - -from decent_array.interoperability._backend_manager import register_backend_listener - -if TYPE_CHECKING: - from decent_array.interoperability._abstracts import Backend - - -_BACKEND_INSTANCE: Backend | None = None -_error = RuntimeError("No backend active: call 'set_backend' with a supported framework to activate one.") - - -def _update_backend(backend: Backend | None) -> None: - global _BACKEND_INSTANCE # noqa: PLW0603 - _BACKEND_INSTANCE = backend - - -register_backend_listener(_update_backend) +from typing import Any + + +_SUPPORTED = { + "bfloat16", + "bool_", + "bytes_", + "complex64", + "complex128", + "complex256", + "float16", + "float32", + "float64", + "float128", + "int8", + "int16", + "int32", + "int64", + "object_", + "qint8", + "qint16", + "qint32", + "quint8", + "quint16", + "uint8", + "uint16", + "uint32", + "uint64", + "unicode_", + "void", +} class dtype: # noqa: N801 """Base class for dtypes.""" def __init__(self, name: str): - if _BACKEND_INSTANCE is None: - raise _error - # name doesn't map to any dtype - if name not in _ALL_DTYPES: - raise ValueError(f"dtype {name} is not supported.") - - # this is None if the dtype is not supported - backend_dtype = getattr(_BACKEND_INSTANCE, name, None) + if name not in _SUPPORTED: + raise ValueError(f"dtype {name} is not supported. Supported dtypes: {', '.join(_SUPPORTED)}") self._name = name - self._backend_dtype: Any = backend_dtype - self._available = backend_dtype is not None + self._backend_dtype: Any = None + self._available = False @property def name(self) -> str: @@ -50,6 +57,11 @@ def available(self) -> bool: """Availability of the dtype (dependent on backend, device, backend settings, OS).""" return self._available + @property + def backend_dtype(self) -> Any: # noqa: ANN401 + """The corresponding backend dtype object.""" + return self._backend_dtype + def __str__(self) -> str: """Name of the dtype.""" return self.name @@ -65,14 +77,8 @@ def __hash__(self) -> int: return hash(self.name) -# TODO create global error message for this when an unavailable dtype is attampted in astype and other funcs -# raise ValueError(f"dtype {name} is not supported.") # TODO add device and backend name to contextualize - -# TODO do we want to have an optional dependence on ml_dtypes to support bfloat16 for numpy? - - _BOOL_DTYPES = { - "bool_": dtype("bool"), + "bool_": dtype("bool_"), } _SIGNED_INT_DTYPES = { @@ -115,9 +121,9 @@ def __hash__(self) -> int: } _MISCELLANEOUS_DTYPES = { - "unicode_": dtype("unicode"), - "bytes_": dtype("bytes"), - "object": dtype("object"), + "unicode_": dtype("unicode_"), + "bytes_": dtype("bytes_"), + "object_": dtype("object_"), "void": dtype("void"), } @@ -130,6 +136,7 @@ def __hash__(self) -> int: _ALL_DTYPES = _BOOL_DTYPES | _NUMERIC_DTYPES | _MISCELLANEOUS_DTYPES + _BACKEND_DTYPE_TO_DTYPE: dict[Any, dtype] = {} for dt in _ALL_DTYPES.values(): if dt.available: From 41aed3b1eb0513254d34e61d6c71e016f85486a1 Mon Sep 17 00:00:00 2001 From: nicola-bastianello Date: Wed, 17 Jun 2026 14:01:37 +0200 Subject: [PATCH 06/23] enh: first working version --- decent_array/_array.py | 19 +--- .../interoperability/_abstracts/backend.py | 4 +- .../interoperability/_backend_manager.py | 15 +++ .../interoperability/_iop/manipulations.py | 4 +- .../interoperability/_jax/jax_backend.py | 28 ++---- .../interoperability/_numpy/numpy_backend.py | 28 ++---- .../_pytorch/pytorch_backend.py | 28 ++---- .../_tensorflow/tensorflow_backend.py | 29 ++---- decent_array/types/__init__.py | 94 ++++++++++++------ decent_array/types/_dtypes.py | 98 +++++++++---------- decent_array/types/_types.py | 37 ------- tests/test_iop_functions.py | 10 +- 12 files changed, 163 insertions(+), 231 deletions(-) diff --git a/decent_array/_array.py b/decent_array/_array.py index 83a36f4..ed9ab13 100644 --- a/decent_array/_array.py +++ b/decent_array/_array.py @@ -29,13 +29,13 @@ from typing import TYPE_CHECKING, Any, Self from decent_array.interoperability._backend_manager import register_backend_listener -from decent_array.types import _STRING_TO_DTYPE +from decent_array.types._dtypes import _BACKEND_DTYPE_TO_DTYPE if TYPE_CHECKING: from numpy.typing import NDArray from decent_array.interoperability._abstracts import Backend - from decent_array.types import ArrayKey, DTypes, ArrayTypes, Devices + from decent_array.types import ArrayKey, ArrayTypes, Devices, dtype _BACKEND_INSTANCE: Backend | None = None @@ -401,19 +401,10 @@ def ndim(self) -> int: return self._backend.ndim(self) @property - def dtype(self) -> DTypes: - """ - Return dtype of the Array as item of DTypes enum. - - Raises: - ValueError: for dtypes that are not supported by all decent-array functions - - """ - # get framework-native dtype as string - # split takes care of types with names like "torch.float32" - dtype_name = str(self.value.dtype).split(".")[-1] + def dtype(self) -> dtype: + """Return dtype of the Array.""" + dtype = _BACKEND_DTYPE_TO_DTYPE.get(self.value.dtype) - dtype = _STRING_TO_DTYPE.get(dtype_name) if dtype is None: raise ValueError(f"dtype {self.value.dtype} is not supported by all decent-array functions.") diff --git a/decent_array/interoperability/_abstracts/backend.py b/decent_array/interoperability/_abstracts/backend.py index 3ca8e0b..be5723f 100644 --- a/decent_array/interoperability/_abstracts/backend.py +++ b/decent_array/interoperability/_abstracts/backend.py @@ -24,7 +24,7 @@ from numpy.typing import NDArray from decent_array import Array - from decent_array.types import ArrayKey, DTypes, ArrayTypes + from decent_array.types import ArrayKey, ArrayTypes, dtype class Backend(ABC): # noqa: PLR0904 @@ -135,7 +135,7 @@ def diagonal(self, x: Array, offset: int = 0) -> Array: """Extract the diagonal entries from a 2-D matrix at the given ``offset``.""" @abstractmethod - def astype(self, x: Array, dtype: DTypes) -> Array: + def astype(self, x: Array, dtype: dtype) -> Array: """Cast ``x`` to a different dtype.""" # Linalg -------------------------------------------------------------- diff --git a/decent_array/interoperability/_backend_manager.py b/decent_array/interoperability/_backend_manager.py index 2c2621c..6e6b74c 100644 --- a/decent_array/interoperability/_backend_manager.py +++ b/decent_array/interoperability/_backend_manager.py @@ -4,7 +4,9 @@ from collections.abc import Callable from contextvars import ContextVar +from decent_array import types from decent_array.types import Devices, Frameworks +from decent_array.types._dtypes import _SUPPORTED from ._abstracts import Backend @@ -69,6 +71,8 @@ def set_backend( for listener in _BACKEND_LISTENERS: listener(_BACKEND_INSTANCE) + _bind_dtypes(_BACKEND_INSTANCE) + def register_backend_listener(listener: Callable[[Backend | None], None]) -> None: """ @@ -202,3 +206,14 @@ def _auto_import(backend: Frameworks) -> None: f"Failed to import the backend module for '{backend.value}'. Ensure the " "corresponding backend package is installed and importable." ) from exc + + +def _bind_dtypes(backend: Backend | None) -> None: + """Bind dtype objects to the corresponding backend dtypes (if available).""" + if backend is None: + return + for name in _SUPPORTED: + dt = getattr(types, name) + backend_dt = getattr(backend, name, None) + dt._available = backend_dt is not None # noqa: SLF001 + dt._backend_dtype = backend_dt # noqa: SLF001 diff --git a/decent_array/interoperability/_iop/manipulations.py b/decent_array/interoperability/_iop/manipulations.py index 1abcd94..0372ccd 100644 --- a/decent_array/interoperability/_iop/manipulations.py +++ b/decent_array/interoperability/_iop/manipulations.py @@ -23,7 +23,7 @@ from decent_array import Array from decent_array.interoperability._abstracts import Backend - from decent_array.types import DTypes, ArrayTypes + from decent_array.types import ArrayTypes, dtype _BACKEND_INSTANCE: Backend | None = None _error = RuntimeError("No backend active: call 'set_backend' with a supported framework to activate one.") @@ -156,7 +156,7 @@ def diagonal(x: Array, offset: int = 0) -> Array: return _BACKEND_INSTANCE.diagonal(x, offset) -def astype(x: Array, dtype: DTypes) -> Array: +def astype(x: Array, dtype: dtype) -> Array: """Cast ``x`` to a different dtype.""" if _BACKEND_INSTANCE is None: raise _error diff --git a/decent_array/interoperability/_jax/jax_backend.py b/decent_array/interoperability/_jax/jax_backend.py index 935be9e..e0ea69c 100644 --- a/decent_array/interoperability/_jax/jax_backend.py +++ b/decent_array/interoperability/_jax/jax_backend.py @@ -22,7 +22,8 @@ from decent_array import Array from decent_array.interoperability._abstracts import Backend from decent_array.interoperability._backend_manager import register_backend -from decent_array.types import ArrayKey, DTypes, ArrayTypes, Devices, Frameworks +from decent_array.types import ArrayKey, ArrayTypes, Devices, Frameworks, dtype +from decent_array.types._dtypes import _ALL_DTYPES def _unwrap(x: Any) -> Any: # noqa: ANN401 @@ -30,23 +31,6 @@ def _unwrap(x: Any) -> Any: # noqa: ANN401 return x.value if type(x) is Array else x -_DTYPE_MAP = { - DTypes.BOOL: jnp.bool_, - DTypes.UINT8: jnp.uint8, - DTypes.UINT16: jnp.uint16, - DTypes.UINT32: jnp.uint32, - DTypes.UINT64: jnp.uint64, - DTypes.INT8: jnp.int8, - DTypes.INT16: jnp.int16, - DTypes.INT32: jnp.int32, - DTypes.INT64: jnp.int64, - DTypes.FLOAT32: jnp.float32, - DTypes.FLOAT64: jnp.float64, - DTypes.COMPLEX64: jnp.complex64, - DTypes.COMPLEX128: jnp.complex128, -} - - class JaxBackend(Backend): # noqa: PLR0904 """JAX implementation of :class:`Backend`.""" @@ -147,10 +131,10 @@ def diagonal(self, x: Array, offset: int = 0) -> Array: raise ValueError(f"diagonal requires a 2-D array, got {x.value.ndim}-D") return Array(jnp.diagonal(x.value, offset=offset)) - def astype(self, x: Array, dtype: DTypes) -> Array: - if dtype not in _DTYPE_MAP: - raise ValueError(f"Unsupported dtype '{dtype.value}' for NumPy backend.") - return Array(jnp.asarray(x.value, dtype=_DTYPE_MAP[dtype])) + def astype(self, x: Array, dtype: dtype) -> Array: + if dtype not in _ALL_DTYPES.values(): + raise ValueError(f"Unsupported dtype '{dtype}' for JAX backend.") + return Array(jnp.asarray(x.value, dtype=dtype.backend_dtype)) # Linalg diff --git a/decent_array/interoperability/_numpy/numpy_backend.py b/decent_array/interoperability/_numpy/numpy_backend.py index 062ee6f..1a40362 100644 --- a/decent_array/interoperability/_numpy/numpy_backend.py +++ b/decent_array/interoperability/_numpy/numpy_backend.py @@ -17,7 +17,8 @@ from decent_array import Array from decent_array.interoperability._abstracts import Backend from decent_array.interoperability._backend_manager import register_backend -from decent_array.types import ArrayKey, DTypes, ArrayTypes, Devices, Frameworks +from decent_array.types import ArrayKey, ArrayTypes, Devices, Frameworks, dtype +from decent_array.types._dtypes import _ALL_DTYPES def _unwrap(x: Any) -> Any: # noqa: ANN401 @@ -31,23 +32,6 @@ def _unwrap(x: Any) -> Any: # noqa: ANN401 return x.value if type(x) is Array else x -_DTYPE_MAP = { - DTypes.BOOL: np.bool_, - DTypes.UINT8: np.uint8, - DTypes.UINT16: np.uint16, - DTypes.UINT32: np.uint32, - DTypes.UINT64: np.uint64, - DTypes.INT8: np.int8, - DTypes.INT16: np.int16, - DTypes.INT32: np.int32, - DTypes.INT64: np.int64, - DTypes.FLOAT32: np.float32, - DTypes.FLOAT64: np.float64, - DTypes.COMPLEX64: np.complex64, - DTypes.COMPLEX128: np.complex128, -} - - class NumpyBackend(Backend): # noqa: PLR0904 """NumPy implementation of :class:`Backend`.""" @@ -148,10 +132,10 @@ def diagonal(self, x: Array, offset: int = 0) -> Array: raise ValueError(f"diagonal requires a 2-D array, got {x.value.ndim}-D") return Array(np.diagonal(x.value, offset=offset)) - def astype(self, x: Array, dtype: DTypes) -> Array: - if dtype not in _DTYPE_MAP: - raise ValueError(f"Unsupported dtype '{dtype.value}' for NumPy backend.") - return Array(np.asarray(x.value, dtype=_DTYPE_MAP[dtype])) + def astype(self, x: Array, dtype: dtype) -> Array: + if dtype not in _ALL_DTYPES.values(): + raise ValueError(f"Unsupported dtype '{dtype}' for NumPy backend.") + return Array(np.asarray(x.value, dtype=dtype.backend_dtype)) # Linalg diff --git a/decent_array/interoperability/_pytorch/pytorch_backend.py b/decent_array/interoperability/_pytorch/pytorch_backend.py index 95f0abd..4397bdb 100644 --- a/decent_array/interoperability/_pytorch/pytorch_backend.py +++ b/decent_array/interoperability/_pytorch/pytorch_backend.py @@ -17,7 +17,8 @@ from decent_array import Array from decent_array.interoperability._abstracts import Backend from decent_array.interoperability._backend_manager import register_backend -from decent_array.types import ArrayKey, DTypes, ArrayTypes, Devices, Frameworks +from decent_array.types import ArrayKey, ArrayTypes, Devices, Frameworks, dtype +from decent_array.types._dtypes import _ALL_DTYPES def _unwrap(x: Any) -> Any: # noqa: ANN401 @@ -25,23 +26,6 @@ def _unwrap(x: Any) -> Any: # noqa: ANN401 return x.value if type(x) is Array else x -_DTYPE_MAP = { - DTypes.BOOL: torch.bool, - DTypes.UINT8: torch.uint8, - DTypes.UINT16: torch.uint16, - DTypes.UINT32: torch.uint32, - DTypes.UINT64: torch.uint64, - DTypes.INT8: torch.int8, - DTypes.INT16: torch.int16, - DTypes.INT32: torch.int32, - DTypes.INT64: torch.int64, - DTypes.FLOAT32: torch.float32, - DTypes.FLOAT64: torch.float64, - DTypes.COMPLEX64: torch.complex64, - DTypes.COMPLEX128: torch.complex128, -} - - class PyTorchBackend(Backend): # noqa: PLR0904 """PyTorch implementation of :class:`Backend`.""" @@ -157,10 +141,10 @@ def diagonal(self, x: Array, offset: int = 0) -> Array: raise ValueError(f"diagonal requires a 2-D array, got {x.value.ndim}-D") return Array(torch.diagonal(x.value, offset=offset)) - def astype(self, x: Array, dtype: DTypes) -> Array: - if dtype not in _DTYPE_MAP: - raise ValueError(f"Unsupported dtype '{dtype.value}' for PyTorch backend.") - return Array(x.value.to(dtype=_DTYPE_MAP[dtype])) + def astype(self, x: Array, dtype: dtype) -> Array: + if dtype not in _ALL_DTYPES.values(): + raise ValueError(f"Unsupported dtype '{dtype}' for PyTorch backend.") + return Array(x.value.to(dtype=dtype.backend_dtype)) # Linalg diff --git a/decent_array/interoperability/_tensorflow/tensorflow_backend.py b/decent_array/interoperability/_tensorflow/tensorflow_backend.py index 2bbb2a4..13a4d31 100644 --- a/decent_array/interoperability/_tensorflow/tensorflow_backend.py +++ b/decent_array/interoperability/_tensorflow/tensorflow_backend.py @@ -20,7 +20,8 @@ from decent_array import Array from decent_array.interoperability._abstracts import Backend from decent_array.interoperability._backend_manager import register_backend -from decent_array.types import ArrayKey, DTypes, ArrayTypes, Devices, Frameworks +from decent_array.types import ArrayKey, ArrayTypes, Devices, Frameworks, dtype +from decent_array.types._dtypes import _ALL_DTYPES def _unwrap(x: Any) -> Any: # noqa: ANN401 @@ -28,24 +29,6 @@ def _unwrap(x: Any) -> Any: # noqa: ANN401 return x.value if type(x) is Array else x -_DTYPE_MAP = { - DTypes.BOOL: tf.bool, - DTypes.UINT8: tf.uint8, - DTypes.INT8: tf.int8, - DTypes.UINT16: tf.uint16, - DTypes.INT16: tf.int16, - DTypes.UINT32: tf.uint32, - DTypes.INT32: tf.int32, - DTypes.UINT64: tf.uint64, - DTypes.INT64: tf.int64, - DTypes.FLOAT16: tf.float16, - DTypes.FLOAT32: tf.float32, - DTypes.FLOAT64: tf.float64, - DTypes.COMPLEX64: tf.complex64, - DTypes.COMPLEX128: tf.complex128, -} - - class TensorflowBackend(Backend): # noqa: PLR0904 """TensorFlow implementation of :class:`Backend`.""" @@ -163,10 +146,10 @@ def diagonal(self, x: Array, offset: int = 0) -> Array: raise ValueError(f"diagonal requires a 2-D tensor, got rank {rank}") return Array(tf.linalg.diag_part(v, k=offset)) - def astype(self, x: Array, dtype: DTypes) -> Array: - if dtype not in _DTYPE_MAP: - raise ValueError(f"Unsupported dtype '{dtype.value}' for TensorFlow backend.") - return Array(tf.cast(x.value, dtype=_DTYPE_MAP[dtype])) + def astype(self, x: Array, dtype: dtype) -> Array: + if dtype not in _ALL_DTYPES.values(): + raise ValueError(f"Unsupported dtype '{dtype}' for TensorFlow backend.") + return Array(tf.cast(x.value, dtype=dtype.backend_dtype)) # Linalg diff --git a/decent_array/types/__init__.py b/decent_array/types/__init__.py index e69df9e..f68d01f 100644 --- a/decent_array/types/__init__.py +++ b/decent_array/types/__init__.py @@ -1,51 +1,83 @@ -from typing import TYPE_CHECKING - -from decent_array.types._dtypes import _ALL_DTYPES, dtype, dtypes -from decent_array.types._types import ( - ArrayKey as ArrayKey, +from decent_array.types._dtypes import ( + bfloat16, + bool_, + bytes_, + complex64, + complex128, + complex256, + dtype, + dtypes, + float16, + float32, + float64, + float128, + int8, + int16, + int32, + int64, + object_, + qint8, + qint16, + qint32, + quint8, + quint16, + uint8, + uint16, + uint32, + uint64, + unicode_, + void, ) from decent_array.types._types import ( - ArrayLike as ArrayLike, -) -from decent_array.types._types import ( - DTypes as DTypes, -) -from decent_array.types._types import ( - ArrayTypes as ArrayTypes, -) -from decent_array.types._types import ( - Devices as Devices, -) -from decent_array.types._types import ( - Frameworks as Frameworks, + ArrayKey, + ArrayLike, + ArrayTypes, + Devices, + Frameworks, ) -_static_exports = [ # noqa: RUF067 +__all_docs__ = [ + "ArrayKey", "ArrayLike", "ArrayTypes", - "ArrayKey", - "Frameworks", "Devices", - "DTypes", + "Frameworks", "dtype", "dtypes", ] -_all_dtypes = list(_ALL_DTYPES.keys()) # noqa: RUF067 -_available_dtypes = [name for name, dt in _ALL_DTYPES.items() if dt.available] # noqa: RUF067 - -__all_docs__ = _static_exports + _all_dtypes - __all__ = [ "ArrayKey", "ArrayLike", - "DTypes", "ArrayTypes", "Devices", "Frameworks", + "bfloat16", + "bool_", + "bytes_", + "complex64", + "complex128", + "complex256", "dtype", "dtypes", + "float16", + "float32", + "float64", + "float128", + "int8", + "int16", + "int32", + "int64", + "object_", + "qint8", + "qint16", + "qint32", + "quint8", + "quint16", + "uint8", + "uint16", + "uint32", + "uint64", + "unicode_", + "void", ] - -if not TYPE_CHECKING: # noqa: RUF067 - __all__ += _available_dtypes # noqa: PLE0605 diff --git a/decent_array/types/_dtypes.py b/decent_array/types/_dtypes.py index eea5074..31dbac0 100644 --- a/decent_array/types/_dtypes.py +++ b/decent_array/types/_dtypes.py @@ -4,7 +4,6 @@ from typing import Any - _SUPPORTED = { "bfloat16", "bool_", @@ -43,6 +42,7 @@ def __init__(self, name: str): if name not in _SUPPORTED: raise ValueError(f"dtype {name} is not supported. Supported dtypes: {', '.join(_SUPPORTED)}") + # initialize with placeholder values self._name = name self._backend_dtype: Any = None self._available = False @@ -77,66 +77,62 @@ def __hash__(self) -> int: return hash(self.name) -_BOOL_DTYPES = { - "bool_": dtype("bool_"), -} +# instantiate all the supported dtypes +bool_ = dtype("bool_") +_BOOL_DTYPES = {"bool_": bool_} -_SIGNED_INT_DTYPES = { - "int8": dtype("int8"), - "int16": dtype("int16"), - "int32": dtype("int32"), - "int64": dtype("int64"), -} +int8 = dtype("int8") +int16 = dtype("int16") +int32 = dtype("int32") +int64 = dtype("int64") +_SIGNED_INT_DTYPES = {"int8": int8, "int16": int16, "int32": int32, "int64": int64} -_UNSIGNED_INT_DTYPES = { - "uint8": dtype("uint8"), - "uint16": dtype("uint16"), - "uint32": dtype("uint32"), - "uint64": dtype("uint64"), -} +uint8 = dtype("uint8") +uint16 = dtype("uint16") +uint32 = dtype("uint32") +uint64 = dtype("uint64") +_UNSIGNED_INT_DTYPES = {"uint8": uint8, "uint16": uint16, "uint32": uint32, "uint64": uint64} +float16 = dtype("float16") +bfloat16 = dtype("bfloat16") +float32 = dtype("float32") +float64 = dtype("float64") +float128 = dtype("float128") _REAL_FLOATING_DTYPES = { - "float16": dtype("float16"), - "bfloat16": dtype("bfloat16"), - "float32": dtype("float32"), - "float64": dtype("float64"), - "float128": dtype("float128"), + "float16": float16, + "bfloat16": bfloat16, + "float32": float32, + "float64": float64, + "float128": float128, } -_COMPLEX_FLOATING_DTYPES = { - "complex64": dtype("complex64"), - "complex128": dtype("complex128"), - "complex256": dtype("complex256"), -} - -_QUANTIZED_SIGNED_INT_DTYPES = { - "qint8": dtype("qint8"), - "qint16": dtype("qint16"), - "qint32": dtype("qint32"), -} - -_QUANTIZED_UNSIGNED_INT_DTYPES = { - "quint8": dtype("quint8"), - "quint16": dtype("quint16"), -} - -_MISCELLANEOUS_DTYPES = { - "unicode_": dtype("unicode_"), - "bytes_": dtype("bytes_"), - "object_": dtype("object_"), - "void": dtype("void"), -} - -_INTEGRAL_DTYPES = (_SIGNED_INT_DTYPES | - _UNSIGNED_INT_DTYPES | - _QUANTIZED_SIGNED_INT_DTYPES | - _QUANTIZED_UNSIGNED_INT_DTYPES - ) +complex64 = dtype("complex64") +complex128 = dtype("complex128") +complex256 = dtype("complex256") +_COMPLEX_FLOATING_DTYPES = {"complex64": complex64, "complex128": complex128, "complex256": complex256} + +qint8 = dtype("qint8") +qint16 = dtype("qint16") +qint32 = dtype("qint32") +_QUANTIZED_SIGNED_INT_DTYPES = {"qint8": qint8, "qint16": qint16, "qint32": qint32} + +quint8 = dtype("quint8") +quint16 = dtype("quint16") +_QUANTIZED_UNSIGNED_INT_DTYPES = {"quint8": quint8, "quint16": quint16} + +unicode_ = dtype("unicode_") +bytes_ = dtype("bytes_") +object_ = dtype("object_") +void = dtype("void") +_MISCELLANEOUS_DTYPES = {"unicode_": unicode_, "bytes_": bytes_, "object_": object_, "void": void} + +_INTEGRAL_DTYPES = ( + _SIGNED_INT_DTYPES | _UNSIGNED_INT_DTYPES | _QUANTIZED_SIGNED_INT_DTYPES | _QUANTIZED_UNSIGNED_INT_DTYPES +) _NUMERIC_DTYPES = _INTEGRAL_DTYPES | _REAL_FLOATING_DTYPES | _COMPLEX_FLOATING_DTYPES _ALL_DTYPES = _BOOL_DTYPES | _NUMERIC_DTYPES | _MISCELLANEOUS_DTYPES - _BACKEND_DTYPE_TO_DTYPE: dict[Any, dtype] = {} for dt in _ALL_DTYPES.values(): if dt.available: diff --git a/decent_array/types/_types.py b/decent_array/types/_types.py index 963cabd..ff5a081 100644 --- a/decent_array/types/_types.py +++ b/decent_array/types/_types.py @@ -5,8 +5,6 @@ from enum import Enum from typing import TYPE_CHECKING, SupportsIndex, TypeAlias, Union -from decent_array.interoperability._backend_manager import register_backend_listener - if TYPE_CHECKING: import jax import numpy @@ -14,19 +12,6 @@ import torch from decent_array._array import Array - from decent_array.interoperability._abstracts import Backend - - -_BACKEND_INSTANCE: Backend | None = None -_error = RuntimeError("No backend active: call 'set_backend' with a supported framework to activate one.") - - -def _update_backend(backend: Backend | None) -> None: - global _BACKEND_INSTANCE # noqa: PLW0603 - _BACKEND_INSTANCE = backend - - -register_backend_listener(_update_backend) ArrayLike: TypeAlias = Union["numpy.ndarray", "torch.Tensor", "tf.Tensor", "jax.Array"] # noqa: UP040 @@ -67,25 +52,3 @@ class Devices(Enum): CPU = "cpu" GPU = "gpu" MPS = "mps" - - -class DTypes(Enum): - """Enum for supported dtypes in decent-array.""" - - BOOL = "bool" - UINT8 = "uint8" - UINT16 = "uint16" - UINT32 = "uint32" - UINT64 = "uint64" - INT8 = "int8" - INT16 = "int16" - INT32 = "int32" - INT64 = "int64" - FLOAT16 = "float16" - FLOAT32 = "float32" - FLOAT64 = "float64" - COMPLEX64 = "complex64" - COMPLEX128 = "complex128" - - -_STRING_TO_DTYPE = {dt.value: dt for dt in DTypes} diff --git a/tests/test_iop_functions.py b/tests/test_iop_functions.py index e4c46ee..b2c74f4 100644 --- a/tests/test_iop_functions.py +++ b/tests/test_iop_functions.py @@ -6,11 +6,11 @@ import pytest import decent_array.interoperability as iop +import decent_array.types as types from decent_array import Array from decent_array.interoperability._backend_manager import reset_backends from decent_array.interoperability._iop.math import iadd, idivide, imultiply, isubtract from decent_array.interoperability._iop.utils import device_to_native, get_item, set_item -from decent_array.types import DTypes def _np(arr: Array) -> np.ndarray: @@ -219,7 +219,7 @@ def test_diagonal_with_offset(backend: tuple) -> None: def test_astype_to_float(backend: tuple) -> None: arr = iop.asarray(3.0) - out = iop.astype(arr, DTypes.FLOAT32) + out = iop.astype(arr, types.float32) np_out = _np(out) assert np_out.dtype == np.float32 assert np_out == pytest.approx(3.0) @@ -227,7 +227,7 @@ def test_astype_to_float(backend: tuple) -> None: def test_astype_to_int(backend: tuple) -> None: arr = iop.asarray(3.0) - out = iop.astype(arr, DTypes.INT32) + out = iop.astype(arr, types.int32) np_out = _np(out) assert np_out.dtype == np.int32 assert int(np_out) == 3 @@ -235,7 +235,7 @@ def test_astype_to_int(backend: tuple) -> None: def test_astype_to_bool(backend: tuple) -> None: arr = iop.asarray(1.0) - out = iop.astype(arr, DTypes.BOOL) + out = iop.astype(arr, types.bool_) np_out = _np(out) assert np_out.dtype == np.bool_ assert bool(np_out) is True @@ -697,7 +697,7 @@ def test_function_raises_when_no_backend() -> None: def test_to_array_round_trip_with_bool(backend: tuple) -> None: arr = iop.asarray(True) - out = iop.astype(arr, DTypes.BOOL) + out = iop.astype(arr, types.bool_) np_out = _np(out) assert np_out.dtype == np.bool_ assert bool(np_out) is True From 86212dafa164518afcf832b36c9dd91b437f3525 Mon Sep 17 00:00:00 2001 From: nicola-bastianello Date: Wed, 17 Jun 2026 17:21:04 +0200 Subject: [PATCH 07/23] enh: improvements --- .../interoperability/_abstracts/backend.py | 5 +++-- .../interoperability/_backend_manager.py | 2 +- decent_array/types/_dtypes.py | 19 +++++++++++++++---- 3 files changed, 19 insertions(+), 7 deletions(-) diff --git a/decent_array/interoperability/_abstracts/backend.py b/decent_array/interoperability/_abstracts/backend.py index be5723f..2507ca5 100644 --- a/decent_array/interoperability/_abstracts/backend.py +++ b/decent_array/interoperability/_abstracts/backend.py @@ -18,13 +18,14 @@ from collections.abc import Sequence from typing import TYPE_CHECKING, Any -from decent_array.types import Devices +from decent_array.types._types import Devices if TYPE_CHECKING: from numpy.typing import NDArray from decent_array import Array - from decent_array.types import ArrayKey, ArrayTypes, dtype + from decent_array.types._dtypes import dtype + from decent_array.types._types import ArrayKey, ArrayTypes class Backend(ABC): # noqa: PLR0904 diff --git a/decent_array/interoperability/_backend_manager.py b/decent_array/interoperability/_backend_manager.py index 6e6b74c..063a5ca 100644 --- a/decent_array/interoperability/_backend_manager.py +++ b/decent_array/interoperability/_backend_manager.py @@ -5,7 +5,7 @@ from contextvars import ContextVar from decent_array import types -from decent_array.types import Devices, Frameworks +from decent_array.types._types import Devices, Frameworks from decent_array.types._dtypes import _SUPPORTED from ._abstracts import Backend diff --git a/decent_array/types/_dtypes.py b/decent_array/types/_dtypes.py index 31dbac0..8b7ef77 100644 --- a/decent_array/types/_dtypes.py +++ b/decent_array/types/_dtypes.py @@ -2,7 +2,16 @@ from __future__ import annotations -from typing import Any +from typing import TYPE_CHECKING, Any + +from decent_array.interoperability import _backend_manager + +if TYPE_CHECKING: + from decent_array.interoperability._abstracts import Backend + + +_error = RuntimeError("No backend active: call 'set_backend' with a supported framework to activate one.") + _SUPPORTED = { "bfloat16", @@ -42,10 +51,12 @@ def __init__(self, name: str): if name not in _SUPPORTED: raise ValueError(f"dtype {name} is not supported. Supported dtypes: {', '.join(_SUPPORTED)}") - # initialize with placeholder values + # initialize with backend dtype; if backend is not initialized, it sets placeholder values and set_backend + # will then bind the backend dtypes + backend_instance = getattr(_backend_manager, "_BACKEND_INSTANCE", None) self._name = name - self._backend_dtype: Any = None - self._available = False + self._backend_dtype: Any = None if backend_instance is None else getattr(backend_instance, name, None) + self._available = self._backend_dtype is not None @property def name(self) -> str: From 249993fd8b1c516969c5cbaceea4f909752e617b Mon Sep 17 00:00:00 2001 From: nicola-bastianello Date: Wed, 17 Jun 2026 17:52:09 +0200 Subject: [PATCH 08/23] enh: added filtering in dtypes function --- decent_array/types/_dtypes.py | 26 ++++++++++++++++++++++++-- 1 file changed, 24 insertions(+), 2 deletions(-) diff --git a/decent_array/types/_dtypes.py b/decent_array/types/_dtypes.py index 8b7ef77..f045752 100644 --- a/decent_array/types/_dtypes.py +++ b/decent_array/types/_dtypes.py @@ -150,6 +150,28 @@ def __hash__(self) -> int: _BACKEND_DTYPE_TO_DTYPE[dt._backend_dtype] = dt # noqa: SLF001 -def dtypes() -> dict[str, dtype]: +_ALIASES = { + "bool": _BOOL_DTYPES, + "signed integer": _SIGNED_INT_DTYPES, + "unsigned integer": _UNSIGNED_INT_DTYPES, + "integral": _INTEGRAL_DTYPES, + "real floating": _REAL_FLOATING_DTYPES, + "complex floating": _COMPLEX_FLOATING_DTYPES, + "numeric": _NUMERIC_DTYPES, +} + + +def dtypes(*, kind: str | tuple[str, ...] | None = None) -> dict[str, dtype]: """Return a dictionary of available dtypes.""" - return {name: dt for name, dt in _ALL_DTYPES.items() if dt.available} + if kind is None: + dtypes = _ALL_DTYPES + elif isinstance(kind, str): + dtypes = _ALIASES.get(kind, {}) + else: + dtypes = { + k: v + for alias in kind if alias in _ALIASES + for k, v in _ALIASES[alias].items() + } + + return {name: dt for name, dt in dtypes.items() if dt.available} From cf6c2609ca0bc93c73beac4dade83a23525497e3 Mon Sep 17 00:00:00 2001 From: nicola-bastianello Date: Wed, 17 Jun 2026 18:22:28 +0200 Subject: [PATCH 09/23] enh: user guide --- decent_array/_array.py | 2 +- .../interoperability/_abstracts/backend.py | 2 +- .../interoperability/_backend_manager.py | 2 +- decent_array/types/_dtypes.py | 12 +-- docs/source/user.rst | 97 ++++++++++++------- 5 files changed, 68 insertions(+), 47 deletions(-) diff --git a/decent_array/_array.py b/decent_array/_array.py index ed9ab13..4866fbc 100644 --- a/decent_array/_array.py +++ b/decent_array/_array.py @@ -402,7 +402,7 @@ def ndim(self) -> int: @property def dtype(self) -> dtype: - """Return dtype of the Array.""" + """Return dtype of the Array.""" # noqa: DOC501 dtype = _BACKEND_DTYPE_TO_DTYPE.get(self.value.dtype) if dtype is None: diff --git a/decent_array/interoperability/_abstracts/backend.py b/decent_array/interoperability/_abstracts/backend.py index 2507ca5..0ca11d8 100644 --- a/decent_array/interoperability/_abstracts/backend.py +++ b/decent_array/interoperability/_abstracts/backend.py @@ -25,7 +25,7 @@ from decent_array import Array from decent_array.types._dtypes import dtype - from decent_array.types._types import ArrayKey, ArrayTypes + from decent_array.types import ArrayKey, ArrayTypes class Backend(ABC): # noqa: PLR0904 diff --git a/decent_array/interoperability/_backend_manager.py b/decent_array/interoperability/_backend_manager.py index 063a5ca..2bbc1e5 100644 --- a/decent_array/interoperability/_backend_manager.py +++ b/decent_array/interoperability/_backend_manager.py @@ -5,8 +5,8 @@ from contextvars import ContextVar from decent_array import types -from decent_array.types._types import Devices, Frameworks from decent_array.types._dtypes import _SUPPORTED +from decent_array.types._types import Devices, Frameworks from ._abstracts import Backend diff --git a/decent_array/types/_dtypes.py b/decent_array/types/_dtypes.py index f045752..e6880a7 100644 --- a/decent_array/types/_dtypes.py +++ b/decent_array/types/_dtypes.py @@ -2,14 +2,10 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Any +from typing import Any from decent_array.interoperability import _backend_manager -if TYPE_CHECKING: - from decent_array.interoperability._abstracts import Backend - - _error = RuntimeError("No backend active: call 'set_backend' with a supported framework to activate one.") @@ -168,10 +164,6 @@ def dtypes(*, kind: str | tuple[str, ...] | None = None) -> dict[str, dtype]: elif isinstance(kind, str): dtypes = _ALIASES.get(kind, {}) else: - dtypes = { - k: v - for alias in kind if alias in _ALIASES - for k, v in _ALIASES[alias].items() - } + dtypes = {k: v for alias in kind if alias in _ALIASES for k, v in _ALIASES[alias].items()} return {name: dt for name, dt in dtypes.items() if dt.available} diff --git a/docs/source/user.rst b/docs/source/user.rst index ad409ba..b9971e2 100644 --- a/docs/source/user.rst +++ b/docs/source/user.rst @@ -13,45 +13,74 @@ Requires `Python 3.13+ `_ dtypes ------ -decent-array exposes a number of dtypes which are bound to the corresponding framework-native dtypes. The following -lists collect the dtypes always supported, and the dtypes supported only by some frameworks/under certain conditions. -The table below reports all the details. +decent-array exposes a number of dtypes which are bound to the corresponding framework-native dtypes. +decent-array exposes dtypes as instances of :class:`~decent_array.types.dtype`, which can be accessed as +`from decent_array.types import float32` or just `types.float32`. Additionally, dtypes can be accessed by +`dtype("float32")`. +dtypes have attributes: :attr:`~decent_array.types.dtype.name`, :attr:`~decent_array.types.dtype.available` (see +discussion below), :attr:`~decent_array.types.dtype.backend_dtype` (which binds the dtype to the corresponding +framework-native dtype). +The following lists the dtypes exposed by decent-array. -dtypes always available -~~~~~~~~~~~~~~~~~~~~~~~ +- Booleans + - `bool_` -- bool -- int8 -- int16 -- int32 -- uint8 -- float16 -- float32 -- complex64 +- Unsigned integers + - `uint8` + - `uint16` + - `uint32` + - `uint64` -dtypes conditionally available -~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +- Signed integers + - `int8` + - `int16` + - `int32` + - `int64` -- int64 -- uint16 -- uint32 -- uint64 -- bfloat16 -- float64 -- float128 -- complex128 -- complex256 -- qint8 -- quint8 -- qint16 -- quint16 -- qint32 -- unicode -- bytes -- object -- void +- Floating point + - `float16` + - `bfloat16` + - `float32` + - `float64` + - `float128` + +- Complex + - `complex64` + - `complex128` + - `complex256` + +- Quantized integers + - `qint8` + - `quint8` + - `qint16` + - `quint16` + - `qint32` + +- Miscellaneous + - `unicode_` + - `bytes_` + - `object_` + - `void` + + + +Availability +~~~~~~~~~~~~ + +Not all dtypes are available in all configurations. The following factors affect dtype availability: +framework, device, OS, framework settings; additionally, some framework-native operations only support a subset of +dtypes. + +There is no reliable way to check which dtypes are available in the current setting, and this is also subject to change +as frameworks develop. In decent-array, we make a best effort to determine whether a dtype is available and, if not, +set :attr:`~decent_array.types.dtype.available`. to `False`. Currently, we mark dtypes as un/available based on the +table below. Additionally, dtypes are marked as unavailable (and :attr:`~decent_array.types.dtype.backend_dtype` is +`None`) if the backend has not been initialized via `set_backend`. + +However, :attr:`~decent_array.types.dtype.available` is generally not a reliable indicator of availability. The most +reliable way is to try operations that involve the dtype and observe if the framework raises an error. .. list-table:: dtype support across frameworks @@ -244,7 +273,7 @@ dtypes conditionally available - ✗ - ✓ - Equivalent to ``np.bytes_`` and ``tf.string`` - * - ``object`` + * - ``object_`` - ✓ - ✗ - ✗ From 8478f20f06fd8e1d87193cb9f7f266b229745917 Mon Sep 17 00:00:00 2001 From: nicola-bastianello Date: Thu, 18 Jun 2026 16:56:42 +0200 Subject: [PATCH 10/23] enh: added constants --- decent_array/__init__.py | 65 ++++++++++++++++++ decent_array/_constants.py | 10 +++ .../interoperability/_abstracts/backend.py | 22 ++++++ .../interoperability/_backend_manager.py | 13 ++++ .../interoperability/_iop/manipulations.py | 3 +- .../interoperability/_jax/jax_backend.py | 26 ++++++- .../interoperability/_numpy/numpy_backend.py | 26 ++++++- .../_pytorch/pytorch_backend.py | 26 ++++++- .../_tensorflow/tensorflow_backend.py | 26 ++++++- decent_array/types/__init__.py | 67 +------------------ decent_array/types/_dtypes.py | 2 - docs/source/user.rst | 10 ++- 12 files changed, 217 insertions(+), 79 deletions(-) create mode 100644 decent_array/_constants.py diff --git a/decent_array/__init__.py b/decent_array/__init__.py index 9893247..e9bd6b4 100644 --- a/decent_array/__init__.py +++ b/decent_array/__init__.py @@ -1,8 +1,73 @@ from decent_array import interoperability, types from decent_array._array import Array +from decent_array._constants import e, inf, nan, pi +from decent_array.types._dtypes import ( + bfloat16, + bool_, + bytes_, + complex64, + complex128, + complex256, + float16, + float32, + float64, + float128, + int8, + int16, + int32, + int64, + object_, + qint8, + qint16, + qint32, + quint8, + quint16, + uint8, + uint16, + uint32, + uint64, + unicode_, + void, +) + +__all_docs__ = [ + "Array", + "interoperability", + "types", +] __all__ = [ "Array", + "bfloat16", + "bool_", + "bytes_", + "complex64", + "complex128", + "complex256", + "e", + "float16", + "float32", + "float64", + "float128", + "inf", + "int8", + "int16", + "int32", + "int64", "interoperability", + "nan", + "object_", + "pi", + "qint8", + "qint16", + "qint32", + "quint8", + "quint16", "types", + "uint8", + "uint16", + "uint32", + "uint64", + "unicode_", + "void", ] diff --git a/decent_array/_constants.py b/decent_array/_constants.py new file mode 100644 index 0000000..a05e32b --- /dev/null +++ b/decent_array/_constants.py @@ -0,0 +1,10 @@ +"""Numerical constants.""" + +import math + +_CONSTANTS = ["e", "inf", "nan", "pi"] + +e = math.e +inf = math.inf +nan = math.nan +pi = math.pi diff --git a/decent_array/interoperability/_abstracts/backend.py b/decent_array/interoperability/_abstracts/backend.py index 0ca11d8..2fed68a 100644 --- a/decent_array/interoperability/_abstracts/backend.py +++ b/decent_array/interoperability/_abstracts/backend.py @@ -526,3 +526,25 @@ def object_(self) -> Any | None: # noqa: ANN401 def void(self) -> Any | None: # noqa: ANN401 """Raw/void dtype (optional).""" return None + + # CONSTANTS ----------------------------------------------------------- + + @property + @abstractmethod + def e(self) -> Any: # noqa: ANN401 + """e = 2.71828...""" # noqa: D403 + + @property + @abstractmethod + def inf(self) -> Any: # noqa: ANN401 + """Infinity.""" + + @property + @abstractmethod + def nan(self) -> Any: # noqa: ANN401 + """Not-a-number.""" + + @property + @abstractmethod + def pi(self) -> Any: # noqa: ANN401 + """pi = 3.14159...""" # noqa: D403 diff --git a/decent_array/interoperability/_backend_manager.py b/decent_array/interoperability/_backend_manager.py index 2bbc1e5..272b798 100644 --- a/decent_array/interoperability/_backend_manager.py +++ b/decent_array/interoperability/_backend_manager.py @@ -4,6 +4,7 @@ from collections.abc import Callable from contextvars import ContextVar +import decent_array._constants as constants from decent_array import types from decent_array.types._dtypes import _SUPPORTED from decent_array.types._types import Devices, Frameworks @@ -72,6 +73,7 @@ def set_backend( listener(_BACKEND_INSTANCE) _bind_dtypes(_BACKEND_INSTANCE) + _bind_constants(_BACKEND_INSTANCE) def register_backend_listener(listener: Callable[[Backend | None], None]) -> None: @@ -217,3 +219,14 @@ def _bind_dtypes(backend: Backend | None) -> None: backend_dt = getattr(backend, name, None) dt._available = backend_dt is not None # noqa: SLF001 dt._backend_dtype = backend_dt # noqa: SLF001 + + +def _bind_constants(backend: Backend | None) -> None: + """Bind constants to the corresponding backend constants.""" + if backend is None: + return + for name in constants._CONSTANTS: # noqa: SLF001 + backend_c = getattr(backend, name, None) + if backend_c is None: + return + setattr(constants, name, backend_c) diff --git a/decent_array/interoperability/_iop/manipulations.py b/decent_array/interoperability/_iop/manipulations.py index 0372ccd..400a51b 100644 --- a/decent_array/interoperability/_iop/manipulations.py +++ b/decent_array/interoperability/_iop/manipulations.py @@ -23,7 +23,8 @@ from decent_array import Array from decent_array.interoperability._abstracts import Backend - from decent_array.types import ArrayTypes, dtype + from decent_array.types._dtypes import dtype + from decent_array.types import ArrayTypes _BACKEND_INSTANCE: Backend | None = None _error = RuntimeError("No backend active: call 'set_backend' with a supported framework to activate one.") diff --git a/decent_array/interoperability/_jax/jax_backend.py b/decent_array/interoperability/_jax/jax_backend.py index e0ea69c..40d9ce4 100644 --- a/decent_array/interoperability/_jax/jax_backend.py +++ b/decent_array/interoperability/_jax/jax_backend.py @@ -22,8 +22,8 @@ from decent_array import Array from decent_array.interoperability._abstracts import Backend from decent_array.interoperability._backend_manager import register_backend -from decent_array.types import ArrayKey, ArrayTypes, Devices, Frameworks, dtype -from decent_array.types._dtypes import _ALL_DTYPES +from decent_array.types._dtypes import _ALL_DTYPES, dtype +from decent_array.types import ArrayKey, ArrayTypes, Devices, Frameworks def _unwrap(x: Any) -> Any: # noqa: ANN401 @@ -435,6 +435,28 @@ def complex128(self) -> Any: # noqa: ANN401 def bfloat16(self) -> Any: # noqa: ANN401 return np.dtype(jnp.bfloat16) + # Constants + + @property + def e(self) -> Any: # noqa: ANN401 + """e = 2.71828...""" # noqa: D403 + return jnp.e + + @property + def inf(self) -> Any: # noqa: ANN401 + """Infinity.""" + return jnp.inf + + @property + def nan(self) -> Any: # noqa: ANN401 + """Not-a-number.""" + return jnp.nan + + @property + def pi(self) -> Any: # noqa: ANN401 + """pi = 3.14159...""" # noqa: D403 + return jnp.pi + _JAX_CONFIG_READ = cast("Callable[[str], Any]", jax.config.read) diff --git a/decent_array/interoperability/_numpy/numpy_backend.py b/decent_array/interoperability/_numpy/numpy_backend.py index 1a40362..c6e31bf 100644 --- a/decent_array/interoperability/_numpy/numpy_backend.py +++ b/decent_array/interoperability/_numpy/numpy_backend.py @@ -17,8 +17,8 @@ from decent_array import Array from decent_array.interoperability._abstracts import Backend from decent_array.interoperability._backend_manager import register_backend -from decent_array.types import ArrayKey, ArrayTypes, Devices, Frameworks, dtype -from decent_array.types._dtypes import _ALL_DTYPES +from decent_array.types._dtypes import _ALL_DTYPES, dtype +from decent_array.types import ArrayKey, ArrayTypes, Devices, Frameworks def _unwrap(x: Any) -> Any: # noqa: ANN401 @@ -443,5 +443,27 @@ def object_(self) -> Any: # noqa: ANN401 def void(self) -> Any: # noqa: ANN401 return np.dtype(np.void) + # Constants + + @property + def e(self) -> Any: # noqa: ANN401 + """e = 2.71828...""" # noqa: D403 + return np.e + + @property + def inf(self) -> Any: # noqa: ANN401 + """Infinity.""" + return np.inf + + @property + def nan(self) -> Any: # noqa: ANN401 + """Not-a-number.""" + return np.nan + + @property + def pi(self) -> Any: # noqa: ANN401 + """pi = 3.14159...""" # noqa: D403 + return np.pi + register_backend(Frameworks.NUMPY, NumpyBackend) diff --git a/decent_array/interoperability/_pytorch/pytorch_backend.py b/decent_array/interoperability/_pytorch/pytorch_backend.py index 4397bdb..3b83458 100644 --- a/decent_array/interoperability/_pytorch/pytorch_backend.py +++ b/decent_array/interoperability/_pytorch/pytorch_backend.py @@ -17,8 +17,8 @@ from decent_array import Array from decent_array.interoperability._abstracts import Backend from decent_array.interoperability._backend_manager import register_backend -from decent_array.types import ArrayKey, ArrayTypes, Devices, Frameworks, dtype -from decent_array.types._dtypes import _ALL_DTYPES +from decent_array.types._dtypes import _ALL_DTYPES, dtype +from decent_array.types import ArrayKey, ArrayTypes, Devices, Frameworks def _unwrap(x: Any) -> Any: # noqa: ANN401 @@ -468,5 +468,27 @@ def quint8(self) -> Any: # noqa: ANN401 def bfloat16(self) -> torch.dtype: return torch.bfloat16 + # Constants + + @property + def e(self) -> Any: # noqa: ANN401 + """e = 2.71828...""" # noqa: D403 + return torch.e + + @property + def inf(self) -> Any: # noqa: ANN401 + """Infinity.""" + return torch.inf + + @property + def nan(self) -> Any: # noqa: ANN401 + """Not-a-number.""" + return torch.nan + + @property + def pi(self) -> Any: # noqa: ANN401 + """pi = 3.14159...""" # noqa: D403 + return torch.pi + register_backend(Frameworks.PYTORCH, PyTorchBackend) diff --git a/decent_array/interoperability/_tensorflow/tensorflow_backend.py b/decent_array/interoperability/_tensorflow/tensorflow_backend.py index 13a4d31..0dc18f8 100644 --- a/decent_array/interoperability/_tensorflow/tensorflow_backend.py +++ b/decent_array/interoperability/_tensorflow/tensorflow_backend.py @@ -20,8 +20,8 @@ from decent_array import Array from decent_array.interoperability._abstracts import Backend from decent_array.interoperability._backend_manager import register_backend -from decent_array.types import ArrayKey, ArrayTypes, Devices, Frameworks, dtype -from decent_array.types._dtypes import _ALL_DTYPES +from decent_array.types._dtypes import _ALL_DTYPES, dtype +from decent_array.types import ArrayKey, ArrayTypes, Devices, Frameworks def _unwrap(x: Any) -> Any: # noqa: ANN401 @@ -499,5 +499,27 @@ def quint16(self) -> tf.dtypes.DType: def bfloat16(self) -> tf.dtypes.DType: return tf.bfloat16 + # Constants + + @property + def e(self) -> Any: # noqa: ANN401 + """e = 2.71828...""" # noqa: D403 + return tf.experimental.numpy.e + + @property + def inf(self) -> Any: # noqa: ANN401 + """Infinity.""" + return tf.experimental.numpy.inf + + @property + def nan(self) -> Any: # noqa: ANN401 + """Not-a-number.""" + return tf.experimental.numpy.nan + + @property + def pi(self) -> Any: # noqa: ANN401 + """pi = 3.14159...""" # noqa: D403 + return tf.experimental.numpy.pi + register_backend(Frameworks.TENSORFLOW, TensorflowBackend) diff --git a/decent_array/types/__init__.py b/decent_array/types/__init__.py index f68d01f..198067c 100644 --- a/decent_array/types/__init__.py +++ b/decent_array/types/__init__.py @@ -1,33 +1,4 @@ -from decent_array.types._dtypes import ( - bfloat16, - bool_, - bytes_, - complex64, - complex128, - complex256, - dtype, - dtypes, - float16, - float32, - float64, - float128, - int8, - int16, - int32, - int64, - object_, - qint8, - qint16, - qint32, - quint8, - quint16, - uint8, - uint16, - uint32, - uint64, - unicode_, - void, -) +from decent_array.types._dtypes import dtype, dtypes from decent_array.types._types import ( ArrayKey, ArrayLike, @@ -36,48 +7,12 @@ Frameworks, ) -__all_docs__ = [ - "ArrayKey", - "ArrayLike", - "ArrayTypes", - "Devices", - "Frameworks", - "dtype", - "dtypes", -] - __all__ = [ "ArrayKey", "ArrayLike", "ArrayTypes", "Devices", "Frameworks", - "bfloat16", - "bool_", - "bytes_", - "complex64", - "complex128", - "complex256", "dtype", "dtypes", - "float16", - "float32", - "float64", - "float128", - "int8", - "int16", - "int32", - "int64", - "object_", - "qint8", - "qint16", - "qint32", - "quint8", - "quint16", - "uint8", - "uint16", - "uint32", - "uint64", - "unicode_", - "void", ] diff --git a/decent_array/types/_dtypes.py b/decent_array/types/_dtypes.py index e6880a7..1d5d3db 100644 --- a/decent_array/types/_dtypes.py +++ b/decent_array/types/_dtypes.py @@ -6,8 +6,6 @@ from decent_array.interoperability import _backend_manager -_error = RuntimeError("No backend active: call 'set_backend' with a supported framework to activate one.") - _SUPPORTED = { "bfloat16", diff --git a/docs/source/user.rst b/docs/source/user.rst index b9971e2..235744e 100644 --- a/docs/source/user.rst +++ b/docs/source/user.rst @@ -10,13 +10,19 @@ Requires `Python 3.13+ `_ pip install decent-array +Constants +--------- + +decent-array exposes `e`, `inf`, `nan`, `pi` constants (`from decent_array import e`), which are bound to the +corresponding framework-native constants. If the constants are not bound, they fall back on the `math` constants. + + dtypes ------ decent-array exposes a number of dtypes which are bound to the corresponding framework-native dtypes. decent-array exposes dtypes as instances of :class:`~decent_array.types.dtype`, which can be accessed as -`from decent_array.types import float32` or just `types.float32`. Additionally, dtypes can be accessed by -`dtype("float32")`. +`from decent_array import float32`. Additionally, dtypes can be accessed by `dtype("float32")`. dtypes have attributes: :attr:`~decent_array.types.dtype.name`, :attr:`~decent_array.types.dtype.available` (see discussion below), :attr:`~decent_array.types.dtype.backend_dtype` (which binds the dtype to the corresponding From 9c189d1e98562ffa1d4eba014f856701e51a52c4 Mon Sep 17 00:00:00 2001 From: nicola-bastianello Date: Thu, 18 Jun 2026 17:25:38 +0200 Subject: [PATCH 11/23] fix: tests and sphinx --- decent_array/interoperability/_abstracts/backend.py | 2 +- decent_array/interoperability/_backend_manager.py | 7 +++---- decent_array/interoperability/_iop/manipulations.py | 2 +- decent_array/interoperability/_jax/jax_backend.py | 2 +- decent_array/interoperability/_numpy/numpy_backend.py | 2 +- .../interoperability/_pytorch/pytorch_backend.py | 2 +- .../interoperability/_tensorflow/tensorflow_backend.py | 2 +- decent_array/types/_dtypes.py | 1 - docs/source/api/decent_array.types.rst | 5 ++++- docs/source/conf.py | 4 ++++ tests/test_iop_functions.py | 10 +++++----- 11 files changed, 22 insertions(+), 17 deletions(-) diff --git a/decent_array/interoperability/_abstracts/backend.py b/decent_array/interoperability/_abstracts/backend.py index 2fed68a..df5e347 100644 --- a/decent_array/interoperability/_abstracts/backend.py +++ b/decent_array/interoperability/_abstracts/backend.py @@ -24,8 +24,8 @@ from numpy.typing import NDArray from decent_array import Array - from decent_array.types._dtypes import dtype from decent_array.types import ArrayKey, ArrayTypes + from decent_array.types._dtypes import dtype class Backend(ABC): # noqa: PLR0904 diff --git a/decent_array/interoperability/_backend_manager.py b/decent_array/interoperability/_backend_manager.py index 272b798..7c30169 100644 --- a/decent_array/interoperability/_backend_manager.py +++ b/decent_array/interoperability/_backend_manager.py @@ -5,8 +5,7 @@ from contextvars import ContextVar import decent_array._constants as constants -from decent_array import types -from decent_array.types._dtypes import _SUPPORTED +import decent_array.types._dtypes as dtypes from decent_array.types._types import Devices, Frameworks from ._abstracts import Backend @@ -214,8 +213,8 @@ def _bind_dtypes(backend: Backend | None) -> None: """Bind dtype objects to the corresponding backend dtypes (if available).""" if backend is None: return - for name in _SUPPORTED: - dt = getattr(types, name) + for name in dtypes._SUPPORTED: # noqa: SLF001 + dt = getattr(dtypes, name) backend_dt = getattr(backend, name, None) dt._available = backend_dt is not None # noqa: SLF001 dt._backend_dtype = backend_dt # noqa: SLF001 diff --git a/decent_array/interoperability/_iop/manipulations.py b/decent_array/interoperability/_iop/manipulations.py index 400a51b..629a2f8 100644 --- a/decent_array/interoperability/_iop/manipulations.py +++ b/decent_array/interoperability/_iop/manipulations.py @@ -23,8 +23,8 @@ from decent_array import Array from decent_array.interoperability._abstracts import Backend - from decent_array.types._dtypes import dtype from decent_array.types import ArrayTypes + from decent_array.types._dtypes import dtype _BACKEND_INSTANCE: Backend | None = None _error = RuntimeError("No backend active: call 'set_backend' with a supported framework to activate one.") diff --git a/decent_array/interoperability/_jax/jax_backend.py b/decent_array/interoperability/_jax/jax_backend.py index 40d9ce4..9b9f9cc 100644 --- a/decent_array/interoperability/_jax/jax_backend.py +++ b/decent_array/interoperability/_jax/jax_backend.py @@ -22,8 +22,8 @@ from decent_array import Array from decent_array.interoperability._abstracts import Backend from decent_array.interoperability._backend_manager import register_backend -from decent_array.types._dtypes import _ALL_DTYPES, dtype from decent_array.types import ArrayKey, ArrayTypes, Devices, Frameworks +from decent_array.types._dtypes import _ALL_DTYPES, dtype def _unwrap(x: Any) -> Any: # noqa: ANN401 diff --git a/decent_array/interoperability/_numpy/numpy_backend.py b/decent_array/interoperability/_numpy/numpy_backend.py index c6e31bf..55e014e 100644 --- a/decent_array/interoperability/_numpy/numpy_backend.py +++ b/decent_array/interoperability/_numpy/numpy_backend.py @@ -17,8 +17,8 @@ from decent_array import Array from decent_array.interoperability._abstracts import Backend from decent_array.interoperability._backend_manager import register_backend -from decent_array.types._dtypes import _ALL_DTYPES, dtype from decent_array.types import ArrayKey, ArrayTypes, Devices, Frameworks +from decent_array.types._dtypes import _ALL_DTYPES, dtype def _unwrap(x: Any) -> Any: # noqa: ANN401 diff --git a/decent_array/interoperability/_pytorch/pytorch_backend.py b/decent_array/interoperability/_pytorch/pytorch_backend.py index 3b83458..daedfa1 100644 --- a/decent_array/interoperability/_pytorch/pytorch_backend.py +++ b/decent_array/interoperability/_pytorch/pytorch_backend.py @@ -17,8 +17,8 @@ from decent_array import Array from decent_array.interoperability._abstracts import Backend from decent_array.interoperability._backend_manager import register_backend -from decent_array.types._dtypes import _ALL_DTYPES, dtype from decent_array.types import ArrayKey, ArrayTypes, Devices, Frameworks +from decent_array.types._dtypes import _ALL_DTYPES, dtype def _unwrap(x: Any) -> Any: # noqa: ANN401 diff --git a/decent_array/interoperability/_tensorflow/tensorflow_backend.py b/decent_array/interoperability/_tensorflow/tensorflow_backend.py index 0dc18f8..f658499 100644 --- a/decent_array/interoperability/_tensorflow/tensorflow_backend.py +++ b/decent_array/interoperability/_tensorflow/tensorflow_backend.py @@ -20,8 +20,8 @@ from decent_array import Array from decent_array.interoperability._abstracts import Backend from decent_array.interoperability._backend_manager import register_backend -from decent_array.types._dtypes import _ALL_DTYPES, dtype from decent_array.types import ArrayKey, ArrayTypes, Devices, Frameworks +from decent_array.types._dtypes import _ALL_DTYPES, dtype def _unwrap(x: Any) -> Any: # noqa: ANN401 diff --git a/decent_array/types/_dtypes.py b/decent_array/types/_dtypes.py index 1d5d3db..545fc36 100644 --- a/decent_array/types/_dtypes.py +++ b/decent_array/types/_dtypes.py @@ -6,7 +6,6 @@ from decent_array.interoperability import _backend_manager - _SUPPORTED = { "bfloat16", "bool_", diff --git a/docs/source/api/decent_array.types.rst b/docs/source/api/decent_array.types.rst index 6a29483..09cdf13 100644 --- a/docs/source/api/decent_array.types.rst +++ b/docs/source/api/decent_array.types.rst @@ -4,4 +4,7 @@ decent\_array.types .. automodule:: decent_array.types :members: :show-inheritance: - :undoc-members: \ No newline at end of file + :undoc-members: + +.. autodata:: decent_array.types.ArrayTypes + :annotation: \ No newline at end of file diff --git a/docs/source/conf.py b/docs/source/conf.py index 4c33da5..4d475a0 100644 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -28,6 +28,10 @@ "sphinx.ext.viewcode", # View source code ] +autodoc_type_aliases = { + "ArrayTypes": "decent_array.types.ArrayTypes", +} + nitpicky = True nitpick_ignore = [ ("py:class", "numpy.float64"), diff --git a/tests/test_iop_functions.py b/tests/test_iop_functions.py index b2c74f4..c115504 100644 --- a/tests/test_iop_functions.py +++ b/tests/test_iop_functions.py @@ -6,7 +6,7 @@ import pytest import decent_array.interoperability as iop -import decent_array.types as types +import decent_array as da from decent_array import Array from decent_array.interoperability._backend_manager import reset_backends from decent_array.interoperability._iop.math import iadd, idivide, imultiply, isubtract @@ -219,7 +219,7 @@ def test_diagonal_with_offset(backend: tuple) -> None: def test_astype_to_float(backend: tuple) -> None: arr = iop.asarray(3.0) - out = iop.astype(arr, types.float32) + out = iop.astype(arr, da.float32) np_out = _np(out) assert np_out.dtype == np.float32 assert np_out == pytest.approx(3.0) @@ -227,7 +227,7 @@ def test_astype_to_float(backend: tuple) -> None: def test_astype_to_int(backend: tuple) -> None: arr = iop.asarray(3.0) - out = iop.astype(arr, types.int32) + out = iop.astype(arr, da.int32) np_out = _np(out) assert np_out.dtype == np.int32 assert int(np_out) == 3 @@ -235,7 +235,7 @@ def test_astype_to_int(backend: tuple) -> None: def test_astype_to_bool(backend: tuple) -> None: arr = iop.asarray(1.0) - out = iop.astype(arr, types.bool_) + out = iop.astype(arr, da.bool_) np_out = _np(out) assert np_out.dtype == np.bool_ assert bool(np_out) is True @@ -697,7 +697,7 @@ def test_function_raises_when_no_backend() -> None: def test_to_array_round_trip_with_bool(backend: tuple) -> None: arr = iop.asarray(True) - out = iop.astype(arr, types.bool_) + out = iop.astype(arr, da.bool_) np_out = _np(out) assert np_out.dtype == np.bool_ assert bool(np_out) is True From 3d02be11c685bd2e0228e0e45de2239fb0642ee9 Mon Sep 17 00:00:00 2001 From: nicola-bastianello Date: Thu, 18 Jun 2026 17:43:50 +0200 Subject: [PATCH 12/23] enh: add test for dtypes --- decent_array/types/_dtypes.py | 2 +- tests/test_backend_manager.py | 19 +++++++++++++++++++ 2 files changed, 20 insertions(+), 1 deletion(-) diff --git a/decent_array/types/_dtypes.py b/decent_array/types/_dtypes.py index 545fc36..2c46c97 100644 --- a/decent_array/types/_dtypes.py +++ b/decent_array/types/_dtypes.py @@ -74,7 +74,7 @@ def __eq__(self, other: object) -> bool: """Check equivalence by ``name`` attributes.""" if not isinstance(other, dtype): return NotImplemented - return self.name == other.name + return self.name == other.name and self.available and other.available def __hash__(self) -> int: """Hash of the dtype.""" diff --git a/tests/test_backend_manager.py b/tests/test_backend_manager.py index 3f79e66..8e890c6 100644 --- a/tests/test_backend_manager.py +++ b/tests/test_backend_manager.py @@ -4,8 +4,10 @@ from typing import TYPE_CHECKING +import numpy as np import pytest +import decent_array._constants as constants from decent_array.interoperability import _backend_manager as backend_manager from decent_array.interoperability._abstracts import Backend from decent_array.interoperability._backend_manager import ( @@ -18,6 +20,7 @@ set_backend, ) from decent_array.types import Devices, Frameworks +from decent_array.types._dtypes import dtype, float32 if TYPE_CHECKING: from collections.abc import Iterator @@ -89,6 +92,22 @@ def test_set_backend_invalid_name_raises() -> None: set_backend("not-a-backend") +def test_set_backend_instantiates_dtypes() -> None: + dt1 = dtype("float32") # backend not set, dtype has placeholder values + set_backend("numpy", "cpu") + dt2 = float32 # global dtype bound during set_dtype + dt3 = dtype("float32") # backend is set, so this is bound to backend dtype + + assert dt1 != dt2 # dtypes are not equal if not avaiable, and dt1 is not available + assert dt2 == dt3 # available because bound to backend, and equal + + dt4 = dtype("int16") + assert dt3 != dt4 + + assert dt2.backend_dtype == np.dtype("float32") # check global dtype is correctly bound to backend + assert dt3.backend_dtype == np.dtype("float32") # check dtypes instantiated after set_backend are correctly bound to backend + + # register_backend ------------------------------------------------------- From f40712f36415fae8c175628e929df5bbf7463c5a Mon Sep 17 00:00:00 2001 From: nicola-bastianello Date: Thu, 18 Jun 2026 19:54:00 +0200 Subject: [PATCH 13/23] enh: add backend name and utils submodule --- decent_array/_array.py | 2 +- decent_array/_utils.py | 14 ++++ .../interoperability/_abstracts/backend.py | 3 +- .../interoperability/_jax/jax_backend.py | 72 ++++++++--------- .../interoperability/_numpy/numpy_backend.py | 78 ++++++++----------- .../_pytorch/pytorch_backend.py | 72 ++++++++--------- .../_tensorflow/tensorflow_backend.py | 72 ++++++++--------- 7 files changed, 153 insertions(+), 160 deletions(-) create mode 100644 decent_array/_utils.py diff --git a/decent_array/_array.py b/decent_array/_array.py index 4866fbc..ca530a0 100644 --- a/decent_array/_array.py +++ b/decent_array/_array.py @@ -406,7 +406,7 @@ def dtype(self) -> dtype: dtype = _BACKEND_DTYPE_TO_DTYPE.get(self.value.dtype) if dtype is None: - raise ValueError(f"dtype {self.value.dtype} is not supported by all decent-array functions.") + raise ValueError(f"dtype {self.value.dtype} is not supported.") return dtype diff --git a/decent_array/_utils.py b/decent_array/_utils.py new file mode 100644 index 0000000..559e3d5 --- /dev/null +++ b/decent_array/_utils.py @@ -0,0 +1,14 @@ +from typing import Any + +from decent_array._array import Array + + +def unwrap(x: Any) -> Any: # noqa: ANN401 + """ + Return the underlying value of an :class:`Array`, or pass ``x`` through. + + Typed as ``Any`` because operator dunders may pass either an :class:`Array` or a + Python scalar; the strict abstract signature would force a ``cast`` at every call + site without runtime benefit. + """ + return x.value if type(x) is Array else x diff --git a/decent_array/interoperability/_abstracts/backend.py b/decent_array/interoperability/_abstracts/backend.py index df5e347..4160f74 100644 --- a/decent_array/interoperability/_abstracts/backend.py +++ b/decent_array/interoperability/_abstracts/backend.py @@ -36,8 +36,9 @@ class Backend(ABC): # noqa: PLR0904 time; that device is the default for all new arrays produced by this backend. """ - def __init__(self, device: Devices = Devices.CPU) -> None: + def __init__(self, device: Devices = Devices.CPU, name: str = "") -> None: self.device: Devices = device + self.name: str = name # Array creation ------------------------------------------------------ diff --git a/decent_array/interoperability/_jax/jax_backend.py b/decent_array/interoperability/_jax/jax_backend.py index 9b9f9cc..a84dd75 100644 --- a/decent_array/interoperability/_jax/jax_backend.py +++ b/decent_array/interoperability/_jax/jax_backend.py @@ -20,22 +20,18 @@ from numpy.typing import NDArray from decent_array import Array +from decent_array._utils import unwrap from decent_array.interoperability._abstracts import Backend from decent_array.interoperability._backend_manager import register_backend from decent_array.types import ArrayKey, ArrayTypes, Devices, Frameworks from decent_array.types._dtypes import _ALL_DTYPES, dtype -def _unwrap(x: Any) -> Any: # noqa: ANN401 - """Return the underlying value of an :class:`Array`, or pass ``x`` through.""" - return x.value if type(x) is Array else x - - class JaxBackend(Backend): # noqa: PLR0904 """JAX implementation of :class:`Backend`.""" def __init__(self, device: Devices = Devices.CPU) -> None: - super().__init__(device) + super().__init__(device, name=Frameworks.JAX.value) self._native_device: jax.Device = self.device_to_native(device) self._key: jax.Array = jax.random.key(time_ns()) @@ -182,52 +178,52 @@ def all(self, x: Array, axis: int | tuple[int, ...] | None = None, keepdims: boo # covers both because PEP 484's numeric tower implicitly admits ``int``. def add(self, x1: int | float | complex | Array, x2: int | float | complex | Array) -> Array: - return Array(jnp.add(_unwrap(x1), _unwrap(x2))) + return Array(jnp.add(unwrap(x1), unwrap(x2))) def iadd[T: Array](self, x1: T, x2: int | float | complex | Array) -> T: - x1.value = jnp.add(x1.value, _unwrap(x2)) + x1.value = jnp.add(x1.value, unwrap(x2)) return x1 def subtract(self, x1: int | float | complex | Array, x2: int | float | complex | Array) -> Array: - return Array(jnp.subtract(_unwrap(x1), _unwrap(x2))) + return Array(jnp.subtract(unwrap(x1), unwrap(x2))) def isubtract[T: Array](self, x1: T, x2: int | float | complex | Array) -> T: - x1.value = jnp.subtract(x1.value, _unwrap(x2)) + x1.value = jnp.subtract(x1.value, unwrap(x2)) return x1 def multiply(self, x1: int | float | complex | Array, x2: int | float | complex | Array) -> Array: - return Array(jnp.multiply(_unwrap(x1), _unwrap(x2))) + return Array(jnp.multiply(unwrap(x1), unwrap(x2))) def imultiply[T: Array](self, x1: T, x2: int | float | complex | Array) -> T: - x1.value = jnp.multiply(x1.value, _unwrap(x2)) + x1.value = jnp.multiply(x1.value, unwrap(x2)) return x1 def divide(self, x1: int | float | complex | Array, x2: int | float | complex | Array) -> Array: - return Array(jnp.divide(_unwrap(x1), _unwrap(x2))) + return Array(jnp.divide(unwrap(x1), unwrap(x2))) def idivide[T: Array](self, x1: T, x2: int | float | complex | Array) -> T: - x1.value = jnp.divide(x1.value, _unwrap(x2)) + x1.value = jnp.divide(x1.value, unwrap(x2)) return x1 def floor_divide(self, x1: int | float | Array, x2: int | float | Array) -> Array: - return Array(jnp.floor_divide(_unwrap(x1), _unwrap(x2))) + return Array(jnp.floor_divide(unwrap(x1), unwrap(x2))) def ifloordiv[T: Array](self, x1: T, x2: int | float | Array) -> T: - x1.value = jnp.floor_divide(x1.value, _unwrap(x2)) + x1.value = jnp.floor_divide(x1.value, unwrap(x2)) return x1 def remainder(self, x1: int | float | Array, x2: int | float | Array) -> Array: - return Array(jnp.remainder(_unwrap(x1), _unwrap(x2))) + return Array(jnp.remainder(unwrap(x1), unwrap(x2))) def imod[T: Array](self, x1: T, x2: int | float | Array) -> T: - x1.value = jnp.remainder(x1.value, _unwrap(x2)) + x1.value = jnp.remainder(x1.value, unwrap(x2)) return x1 def pow(self, x1: int | float | complex | Array, x2: int | float | complex | Array) -> Array: - return Array(jnp.power(_unwrap(x1), _unwrap(x2))) + return Array(jnp.power(unwrap(x1), unwrap(x2))) def ipow[T: Array](self, x1: T, x2: int | float | complex | Array) -> T: - x1.value = jnp.power(x1.value, _unwrap(x2)) + x1.value = jnp.power(x1.value, unwrap(x2)) return x1 def negative(self, x: Array) -> Array: @@ -242,61 +238,61 @@ def sqrt(self, x: Array) -> Array: # Comparisons def equal(self, x1: int | float | complex | Array, x2: int | float | complex | Array) -> Array: - return Array(jnp.equal(_unwrap(x1), _unwrap(x2))) + return Array(jnp.equal(unwrap(x1), unwrap(x2))) def not_equal(self, x1: int | float | complex | Array, x2: int | float | complex | Array) -> Array: - return Array(jnp.not_equal(_unwrap(x1), _unwrap(x2))) + return Array(jnp.not_equal(unwrap(x1), unwrap(x2))) def less(self, x1: int | float | complex | Array, x2: int | float | complex | Array) -> Array: - return Array(jnp.less(_unwrap(x1), _unwrap(x2))) + return Array(jnp.less(unwrap(x1), unwrap(x2))) def less_equal(self, x1: int | float | complex | Array, x2: int | float | complex | Array) -> Array: - return Array(jnp.less_equal(_unwrap(x1), _unwrap(x2))) + return Array(jnp.less_equal(unwrap(x1), unwrap(x2))) def greater(self, x1: int | float | complex | Array, x2: int | float | complex | Array) -> Array: - return Array(jnp.greater(_unwrap(x1), _unwrap(x2))) + return Array(jnp.greater(unwrap(x1), unwrap(x2))) def greater_equal(self, x1: int | float | complex | Array, x2: int | float | complex | Array) -> Array: - return Array(jnp.greater_equal(_unwrap(x1), _unwrap(x2))) + return Array(jnp.greater_equal(unwrap(x1), unwrap(x2))) # Bitwise def bitwise_and(self, x1: bool | int | Array, x2: bool | int | Array) -> Array: - return Array(jnp.bitwise_and(_unwrap(x1), _unwrap(x2))) + return Array(jnp.bitwise_and(unwrap(x1), unwrap(x2))) def iand[T: Array](self, x1: T, x2: bool | int | Array) -> T: - x1.value = jnp.bitwise_and(x1.value, _unwrap(x2)) + x1.value = jnp.bitwise_and(x1.value, unwrap(x2)) return x1 def bitwise_invert(self, x: Array) -> Array: return Array(jnp.bitwise_not(x.value)) def bitwise_or(self, x1: bool | int | Array, x2: bool | int | Array) -> Array: - return Array(jnp.bitwise_or(_unwrap(x1), _unwrap(x2))) + return Array(jnp.bitwise_or(unwrap(x1), unwrap(x2))) def ior[T: Array](self, x1: T, x2: bool | int | Array) -> T: - x1.value = jnp.bitwise_or(x1.value, _unwrap(x2)) + x1.value = jnp.bitwise_or(x1.value, unwrap(x2)) return x1 def bitwise_xor(self, x1: bool | int | Array, x2: bool | int | Array) -> Array: - return Array(jnp.bitwise_xor(_unwrap(x1), _unwrap(x2))) + return Array(jnp.bitwise_xor(unwrap(x1), unwrap(x2))) def ixor[T: Array](self, x1: T, x2: bool | int | Array) -> T: - x1.value = jnp.bitwise_xor(x1.value, _unwrap(x2)) + x1.value = jnp.bitwise_xor(x1.value, unwrap(x2)) return x1 def bitwise_left_shift(self, x1: int | Array, x2: int | Array) -> Array: - return Array(jnp.left_shift(_unwrap(x1), _unwrap(x2))) + return Array(jnp.left_shift(unwrap(x1), unwrap(x2))) def ilshift[T: Array](self, x1: T, x2: int | Array) -> T: - x1.value = jnp.left_shift(x1.value, _unwrap(x2)) + x1.value = jnp.left_shift(x1.value, unwrap(x2)) return x1 def bitwise_right_shift(self, x1: int | Array, x2: int | Array) -> Array: - return Array(jnp.right_shift(_unwrap(x1), _unwrap(x2))) + return Array(jnp.right_shift(unwrap(x1), unwrap(x2))) def irshift[T: Array](self, x1: T, x2: int | Array) -> T: - x1.value = jnp.right_shift(x1.value, _unwrap(x2)) + x1.value = jnp.right_shift(x1.value, unwrap(x2)) return x1 # Operators @@ -305,7 +301,7 @@ def sign(self, x: Array) -> Array: return Array(jnp.sign(x.value)) def maximum(self, x1: int | float | complex | Array, x2: int | float | complex | Array) -> Array: - return Array(jnp.maximum(_unwrap(x1), _unwrap(x2))) + return Array(jnp.maximum(unwrap(x1), unwrap(x2))) def argmax(self, x: Array, axis: int | None = None, keepdims: bool = False) -> Array: return Array(jnp.argmax(x.value, axis=axis, keepdims=keepdims)) @@ -315,7 +311,7 @@ def argmin(self, x: Array, axis: int | None = None, keepdims: bool = False) -> A def set_item(self, x: Array, key: ArrayKey, value: bool | int | float | complex | Array) -> None: # JAX arrays are immutable; rebind the wrapper to a new array with `key` updated. - x.value = x.value.at[key].set(_unwrap(value)) + x.value = x.value.at[key].set(unwrap(value)) def get_item(self, x: Array, key: ArrayKey) -> Array: return Array(x.value[key]) diff --git a/decent_array/interoperability/_numpy/numpy_backend.py b/decent_array/interoperability/_numpy/numpy_backend.py index 55e014e..d94727f 100644 --- a/decent_array/interoperability/_numpy/numpy_backend.py +++ b/decent_array/interoperability/_numpy/numpy_backend.py @@ -15,30 +15,20 @@ from numpy.typing import NDArray from decent_array import Array +from decent_array._utils import unwrap from decent_array.interoperability._abstracts import Backend from decent_array.interoperability._backend_manager import register_backend from decent_array.types import ArrayKey, ArrayTypes, Devices, Frameworks from decent_array.types._dtypes import _ALL_DTYPES, dtype -def _unwrap(x: Any) -> Any: # noqa: ANN401 - """ - Return the underlying value of an :class:`Array`, or pass ``x`` through. - - Typed as ``Any`` because operator dunders may pass either an :class:`Array` or a - Python scalar; the strict abstract signature would force a ``cast`` at every call - site without runtime benefit. - """ - return x.value if type(x) is Array else x - - class NumpyBackend(Backend): # noqa: PLR0904 """NumPy implementation of :class:`Backend`.""" def __init__(self, device: Devices = Devices.CPU) -> None: if device != Devices.CPU: raise ValueError(f"NumPy backend only supports CPU, got '{device.value}'.") - super().__init__(device) + super().__init__(device, name=Frameworks.NUMPY.value) self._rng: np.random.Generator = np.random.default_rng() # Array creation @@ -182,52 +172,52 @@ def all(self, x: Array, axis: int | tuple[int, ...] | None = None, keepdims: boo # ``Array | float`` covers both: PEP 484's numeric tower implicitly admits ``int``. def add(self, x1: int | float | complex | Array, x2: int | float | complex | Array) -> Array: - return Array(np.add(_unwrap(x1), _unwrap(x2))) + return Array(np.add(unwrap(x1), unwrap(x2))) def iadd[T: Array](self, x1: T, x2: int | float | complex | Array) -> T: - x1.value += _unwrap(x2) + x1.value += unwrap(x2) return x1 def subtract(self, x1: int | float | complex | Array, x2: int | float | complex | Array) -> Array: - return Array(np.subtract(_unwrap(x1), _unwrap(x2))) + return Array(np.subtract(unwrap(x1), unwrap(x2))) def isubtract[T: Array](self, x1: T, x2: int | float | complex | Array) -> T: - x1.value -= _unwrap(x2) + x1.value -= unwrap(x2) return x1 def multiply(self, x1: int | float | complex | Array, x2: int | float | complex | Array) -> Array: - return Array(np.multiply(_unwrap(x1), _unwrap(x2))) + return Array(np.multiply(unwrap(x1), unwrap(x2))) def imultiply[T: Array](self, x1: T, x2: int | float | complex | Array) -> T: - x1.value *= _unwrap(x2) + x1.value *= unwrap(x2) return x1 def divide(self, x1: int | float | complex | Array, x2: int | float | complex | Array) -> Array: - return Array(np.divide(_unwrap(x1), _unwrap(x2))) + return Array(np.divide(unwrap(x1), unwrap(x2))) def idivide[T: Array](self, x1: T, x2: int | float | complex | Array) -> T: - x1.value /= _unwrap(x2) + x1.value /= unwrap(x2) return x1 def floor_divide(self, x1: int | float | Array, x2: int | float | Array) -> Array: - return Array(np.floor_divide(_unwrap(x1), _unwrap(x2))) + return Array(np.floor_divide(unwrap(x1), unwrap(x2))) def ifloordiv[T: Array](self, x1: T, x2: int | float | Array) -> T: - x1.value //= _unwrap(x2) + x1.value //= unwrap(x2) return x1 def remainder(self, x1: int | float | Array, x2: int | float | Array) -> Array: - return Array(np.remainder(_unwrap(x1), _unwrap(x2))) + return Array(np.remainder(unwrap(x1), unwrap(x2))) def imod[T: Array](self, x1: T, x2: int | float | Array) -> T: - x1.value %= _unwrap(x2) + x1.value %= unwrap(x2) return x1 def pow(self, x1: int | float | complex | Array, x2: int | float | complex | Array) -> Array: - return Array(np.power(_unwrap(x1), _unwrap(x2))) + return Array(np.power(unwrap(x1), unwrap(x2))) def ipow[T: Array](self, x1: T, x2: int | float | complex | Array) -> T: - x1.value **= _unwrap(x2) + x1.value **= unwrap(x2) return x1 def negative(self, x: Array) -> Array: @@ -242,61 +232,61 @@ def sqrt(self, x: Array) -> Array: # Comparisons def equal(self, x1: int | float | complex | Array, x2: int | float | complex | Array) -> Array: - return Array(np.equal(_unwrap(x1), _unwrap(x2))) + return Array(np.equal(unwrap(x1), unwrap(x2))) def not_equal(self, x1: int | float | complex | Array, x2: int | float | complex | Array) -> Array: - return Array(np.not_equal(_unwrap(x1), _unwrap(x2))) + return Array(np.not_equal(unwrap(x1), unwrap(x2))) def less(self, x1: int | float | complex | Array, x2: int | float | complex | Array) -> Array: - return Array(np.less(_unwrap(x1), _unwrap(x2))) + return Array(np.less(unwrap(x1), unwrap(x2))) def less_equal(self, x1: int | float | complex | Array, x2: int | float | complex | Array) -> Array: - return Array(np.less_equal(_unwrap(x1), _unwrap(x2))) + return Array(np.less_equal(unwrap(x1), unwrap(x2))) def greater(self, x1: int | float | complex | Array, x2: int | float | complex | Array) -> Array: - return Array(np.greater(_unwrap(x1), _unwrap(x2))) + return Array(np.greater(unwrap(x1), unwrap(x2))) def greater_equal(self, x1: int | float | complex | Array, x2: int | float | complex | Array) -> Array: - return Array(np.greater_equal(_unwrap(x1), _unwrap(x2))) + return Array(np.greater_equal(unwrap(x1), unwrap(x2))) # Bitwise def bitwise_and(self, x1: bool | int | Array, x2: bool | int | Array) -> Array: - return Array(np.bitwise_and(_unwrap(x1), _unwrap(x2))) + return Array(np.bitwise_and(unwrap(x1), unwrap(x2))) def iand[T: Array](self, x1: T, x2: bool | int | Array) -> T: - x1.value &= _unwrap(x2) + x1.value &= unwrap(x2) return x1 def bitwise_invert(self, x: Array) -> Array: return Array(np.bitwise_not(x.value)) def bitwise_or(self, x1: bool | int | Array, x2: bool | int | Array) -> Array: - return Array(np.bitwise_or(_unwrap(x1), _unwrap(x2))) + return Array(np.bitwise_or(unwrap(x1), unwrap(x2))) def ior[T: Array](self, x1: T, x2: bool | int | Array) -> T: - x1.value |= _unwrap(x2) + x1.value |= unwrap(x2) return x1 def bitwise_xor(self, x1: bool | int | Array, x2: bool | int | Array) -> Array: - return Array(np.bitwise_xor(_unwrap(x1), _unwrap(x2))) + return Array(np.bitwise_xor(unwrap(x1), unwrap(x2))) def ixor[T: Array](self, x1: T, x2: bool | int | Array) -> T: - x1.value ^= _unwrap(x2) + x1.value ^= unwrap(x2) return x1 def bitwise_left_shift(self, x1: int | Array, x2: int | Array) -> Array: - return Array(np.left_shift(_unwrap(x1), _unwrap(x2))) + return Array(np.left_shift(unwrap(x1), unwrap(x2))) def ilshift[T: Array](self, x1: T, x2: int | Array) -> T: - x1.value <<= _unwrap(x2) + x1.value <<= unwrap(x2) return x1 def bitwise_right_shift(self, x1: int | Array, x2: int | Array) -> Array: - return Array(np.right_shift(_unwrap(x1), _unwrap(x2))) + return Array(np.right_shift(unwrap(x1), unwrap(x2))) def irshift[T: Array](self, x1: T, x2: int | Array) -> T: - x1.value >>= _unwrap(x2) + x1.value >>= unwrap(x2) return x1 # Operators @@ -305,7 +295,7 @@ def sign(self, x: Array) -> Array: return Array(np.sign(x.value)) def maximum(self, x1: int | float | complex | Array, x2: int | float | complex | Array) -> Array: - return Array(np.maximum(_unwrap(x1), _unwrap(x2))) + return Array(np.maximum(unwrap(x1), unwrap(x2))) def argmax(self, x: Array, axis: int | None = None, keepdims: bool = False) -> Array: return Array(np.argmax(x.value, axis=axis, keepdims=keepdims)) @@ -314,7 +304,7 @@ def argmin(self, x: Array, axis: int | None = None, keepdims: bool = False) -> A return Array(np.argmin(x.value, axis=axis, keepdims=keepdims)) def set_item(self, x: Array, key: ArrayKey, value: bool | int | float | complex | Array) -> None: - x.value[key] = _unwrap(value) + x.value[key] = unwrap(value) def get_item(self, x: Array, key: ArrayKey) -> Array: return Array(x.value[key]) diff --git a/decent_array/interoperability/_pytorch/pytorch_backend.py b/decent_array/interoperability/_pytorch/pytorch_backend.py index daedfa1..2936944 100644 --- a/decent_array/interoperability/_pytorch/pytorch_backend.py +++ b/decent_array/interoperability/_pytorch/pytorch_backend.py @@ -15,22 +15,18 @@ from numpy.typing import NDArray from decent_array import Array +from decent_array._utils import unwrap from decent_array.interoperability._abstracts import Backend from decent_array.interoperability._backend_manager import register_backend from decent_array.types import ArrayKey, ArrayTypes, Devices, Frameworks from decent_array.types._dtypes import _ALL_DTYPES, dtype -def _unwrap(x: Any) -> Any: # noqa: ANN401 - """Return the underlying value of an :class:`Array`, or pass ``x`` through.""" - return x.value if type(x) is Array else x - - class PyTorchBackend(Backend): # noqa: PLR0904 """PyTorch implementation of :class:`Backend`.""" def __init__(self, device: Devices = Devices.CPU) -> None: - super().__init__(device) + super().__init__(device, name=Frameworks.PYTORCH.value) self._native_device: str = self.device_to_native(device) self._generator: torch.Generator = torch.Generator(device=self._native_device) @@ -203,52 +199,52 @@ def all(self, x: Array, axis: int | tuple[int, ...] | None = None, keepdims: boo # ``Array | float`` covers both: PEP 484's numeric tower implicitly admits ``int``. def add(self, x1: int | float | complex | Array, x2: int | float | complex | Array) -> Array: - return Array(torch.add(_unwrap(x1), _unwrap(x2))) + return Array(torch.add(unwrap(x1), unwrap(x2))) def iadd[T: Array](self, x1: T, x2: int | float | complex | Array) -> T: - x1.value.add_(_unwrap(x2)) + x1.value.add_(unwrap(x2)) return x1 def subtract(self, x1: int | float | complex | Array, x2: int | float | complex | Array) -> Array: - return Array(torch.sub(_unwrap(x1), _unwrap(x2))) + return Array(torch.sub(unwrap(x1), unwrap(x2))) def isubtract[T: Array](self, x1: T, x2: int | float | complex | Array) -> T: - x1.value.sub_(_unwrap(x2)) + x1.value.sub_(unwrap(x2)) return x1 def multiply(self, x1: int | float | complex | Array, x2: int | float | complex | Array) -> Array: - return Array(torch.mul(_unwrap(x1), _unwrap(x2))) + return Array(torch.mul(unwrap(x1), unwrap(x2))) def imultiply[T: Array](self, x1: T, x2: int | float | complex | Array) -> T: - x1.value.mul_(_unwrap(x2)) + x1.value.mul_(unwrap(x2)) return x1 def divide(self, x1: int | float | complex | Array, x2: int | float | complex | Array) -> Array: - return Array(torch.div(_unwrap(x1), _unwrap(x2))) + return Array(torch.div(unwrap(x1), unwrap(x2))) def idivide[T: Array](self, x1: T, x2: int | float | complex | Array) -> T: - x1.value.div_(_unwrap(x2)) + x1.value.div_(unwrap(x2)) return x1 def floor_divide(self, x1: int | float | Array, x2: int | float | Array) -> Array: - return Array(torch.floor_divide(_unwrap(x1), _unwrap(x2))) + return Array(torch.floor_divide(unwrap(x1), unwrap(x2))) def ifloordiv[T: Array](self, x1: T, x2: int | float | Array) -> T: - x1.value.floor_divide_(_unwrap(x2)) + x1.value.floor_divide_(unwrap(x2)) return x1 def remainder(self, x1: int | float | Array, x2: int | float | Array) -> Array: - return Array(torch.remainder(_unwrap(x1), _unwrap(x2))) + return Array(torch.remainder(unwrap(x1), unwrap(x2))) def imod[T: Array](self, x1: T, x2: int | float | Array) -> T: - x1.value.remainder_(_unwrap(x2)) + x1.value.remainder_(unwrap(x2)) return x1 def pow(self, x1: int | float | complex | Array, x2: int | float | complex | Array) -> Array: - return Array(torch.pow(_unwrap(x1), _unwrap(x2))) + return Array(torch.pow(unwrap(x1), unwrap(x2))) def ipow[T: Array](self, x1: T, x2: int | float | complex | Array) -> T: - x1.value.pow_(_unwrap(x2)) + x1.value.pow_(unwrap(x2)) return x1 def negative(self, x: Array) -> Array: @@ -263,61 +259,61 @@ def sqrt(self, x: Array) -> Array: # Comparisons def equal(self, x1: int | float | complex | Array, x2: int | float | complex | Array) -> Array: - return Array(torch.eq(_unwrap(x1), _unwrap(x2))) + return Array(torch.eq(unwrap(x1), unwrap(x2))) def not_equal(self, x1: int | float | complex | Array, x2: int | float | complex | Array) -> Array: - return Array(torch.ne(_unwrap(x1), _unwrap(x2))) + return Array(torch.ne(unwrap(x1), unwrap(x2))) def less(self, x1: int | float | complex | Array, x2: int | float | complex | Array) -> Array: - return Array(torch.lt(_unwrap(x1), _unwrap(x2))) + return Array(torch.lt(unwrap(x1), unwrap(x2))) def less_equal(self, x1: int | float | complex | Array, x2: int | float | complex | Array) -> Array: - return Array(torch.le(_unwrap(x1), _unwrap(x2))) + return Array(torch.le(unwrap(x1), unwrap(x2))) def greater(self, x1: int | float | complex | Array, x2: int | float | complex | Array) -> Array: - return Array(torch.gt(_unwrap(x1), _unwrap(x2))) + return Array(torch.gt(unwrap(x1), unwrap(x2))) def greater_equal(self, x1: int | float | complex | Array, x2: int | float | complex | Array) -> Array: - return Array(torch.ge(_unwrap(x1), _unwrap(x2))) + return Array(torch.ge(unwrap(x1), unwrap(x2))) # Bitwise def bitwise_and(self, x1: bool | int | Array, x2: bool | int | Array) -> Array: - return Array(torch.bitwise_and(_unwrap(x1), _unwrap(x2))) + return Array(torch.bitwise_and(unwrap(x1), unwrap(x2))) def iand[T: Array](self, x1: T, x2: bool | int | Array) -> T: - x1.value.bitwise_and_(_unwrap(x2)) + x1.value.bitwise_and_(unwrap(x2)) return x1 def bitwise_invert(self, x: Array) -> Array: return Array(torch.bitwise_not(x.value)) def bitwise_or(self, x1: bool | int | Array, x2: bool | int | Array) -> Array: - return Array(torch.bitwise_or(_unwrap(x1), _unwrap(x2))) + return Array(torch.bitwise_or(unwrap(x1), unwrap(x2))) def ior[T: Array](self, x1: T, x2: bool | int | Array) -> T: - x1.value.bitwise_or_(_unwrap(x2)) + x1.value.bitwise_or_(unwrap(x2)) return x1 def bitwise_xor(self, x1: bool | int | Array, x2: bool | int | Array) -> Array: - return Array(torch.bitwise_xor(_unwrap(x1), _unwrap(x2))) + return Array(torch.bitwise_xor(unwrap(x1), unwrap(x2))) def ixor[T: Array](self, x1: T, x2: bool | int | Array) -> T: - x1.value.bitwise_xor_(_unwrap(x2)) + x1.value.bitwise_xor_(unwrap(x2)) return x1 def bitwise_left_shift(self, x1: int | Array, x2: int | Array) -> Array: - return Array(torch.bitwise_left_shift(_unwrap(x1), _unwrap(x2))) + return Array(torch.bitwise_left_shift(unwrap(x1), unwrap(x2))) def ilshift[T: Array](self, x1: T, x2: int | Array) -> T: - x1.value.bitwise_left_shift_(_unwrap(x2)) + x1.value.bitwise_left_shift_(unwrap(x2)) return x1 def bitwise_right_shift(self, x1: int | Array, x2: int | Array) -> Array: - return Array(torch.bitwise_right_shift(_unwrap(x1), _unwrap(x2))) + return Array(torch.bitwise_right_shift(unwrap(x1), unwrap(x2))) def irshift[T: Array](self, x1: T, x2: int | Array) -> T: - x1.value.bitwise_right_shift_(_unwrap(x2)) + x1.value.bitwise_right_shift_(unwrap(x2)) return x1 # Operators @@ -326,7 +322,7 @@ def sign(self, x: Array) -> Array: return Array(torch.sign(x.value)) def maximum(self, x1: int | float | complex | Array, x2: int | float | complex | Array) -> Array: - a, b = _unwrap(x1), _unwrap(x2) + a, b = unwrap(x1), unwrap(x2) # torch.maximum requires both operands to be Tensors; lift Python scalars to # match the dtype/device of the tensor operand so the contract matches numpy. if not isinstance(a, torch.Tensor): @@ -343,7 +339,7 @@ def argmin(self, x: Array, axis: int | None = None, keepdims: bool = False) -> A return Array(torch.argmin(x.value, dim=axis, keepdim=keepdims)) def set_item(self, x: Array, key: ArrayKey, value: bool | int | float | complex | Array) -> None: - x.value[key] = _unwrap(value) + x.value[key] = unwrap(value) def get_item(self, x: Array, key: ArrayKey) -> Array: return Array(x.value[key]) diff --git a/decent_array/interoperability/_tensorflow/tensorflow_backend.py b/decent_array/interoperability/_tensorflow/tensorflow_backend.py index f658499..498b9d9 100644 --- a/decent_array/interoperability/_tensorflow/tensorflow_backend.py +++ b/decent_array/interoperability/_tensorflow/tensorflow_backend.py @@ -18,22 +18,18 @@ from numpy.typing import NDArray from decent_array import Array +from decent_array._utils import unwrap from decent_array.interoperability._abstracts import Backend from decent_array.interoperability._backend_manager import register_backend from decent_array.types import ArrayKey, ArrayTypes, Devices, Frameworks from decent_array.types._dtypes import _ALL_DTYPES, dtype -def _unwrap(x: Any) -> Any: # noqa: ANN401 - """Return the underlying value of an :class:`Array`, or pass ``x`` through.""" - return x.value if type(x) is Array else x - - class TensorflowBackend(Backend): # noqa: PLR0904 """TensorFlow implementation of :class:`Backend`.""" def __init__(self, device: Devices = Devices.CPU) -> None: - super().__init__(device) + super().__init__(device, name=Frameworks.TENSORFLOW.value) self._native_device: str = self.device_to_native(device) self._generator: tf.random.Generator = tf.random.Generator.from_non_deterministic_state(alg="philox") @@ -210,52 +206,52 @@ def all(self, x: Array, axis: int | tuple[int, ...] | None = None, keepdims: boo # covers both because PEP 484's numeric tower implicitly admits ``int``. def add(self, x1: int | float | complex | Array, x2: int | float | complex | Array) -> Array: - return Array(tf.add(_unwrap(x1), _unwrap(x2))) + return Array(tf.add(unwrap(x1), unwrap(x2))) def iadd[T: Array](self, x1: T, x2: int | float | complex | Array) -> T: - x1.value = tf.add(x1.value, _unwrap(x2)) + x1.value = tf.add(x1.value, unwrap(x2)) return x1 def subtract(self, x1: int | float | complex | Array, x2: int | float | complex | Array) -> Array: - return Array(tf.subtract(_unwrap(x1), _unwrap(x2))) + return Array(tf.subtract(unwrap(x1), unwrap(x2))) def isubtract[T: Array](self, x1: T, x2: int | float | complex | Array) -> T: - x1.value = tf.subtract(x1.value, _unwrap(x2)) + x1.value = tf.subtract(x1.value, unwrap(x2)) return x1 def multiply(self, x1: int | float | complex | Array, x2: int | float | complex | Array) -> Array: - return Array(tf.multiply(_unwrap(x1), _unwrap(x2))) + return Array(tf.multiply(unwrap(x1), unwrap(x2))) def imultiply[T: Array](self, x1: T, x2: int | float | complex | Array) -> T: - x1.value = tf.multiply(x1.value, _unwrap(x2)) + x1.value = tf.multiply(x1.value, unwrap(x2)) return x1 def divide(self, x1: int | float | complex | Array, x2: int | float | complex | Array) -> Array: - return Array(tf.divide(_unwrap(x1), _unwrap(x2))) + return Array(tf.divide(unwrap(x1), unwrap(x2))) def idivide[T: Array](self, x1: T, x2: int | float | complex | Array) -> T: - x1.value = tf.divide(x1.value, _unwrap(x2)) + x1.value = tf.divide(x1.value, unwrap(x2)) return x1 def floor_divide(self, x1: int | float | Array, x2: int | float | Array) -> Array: - return Array(tf.math.floordiv(_unwrap(x1), _unwrap(x2))) + return Array(tf.math.floordiv(unwrap(x1), unwrap(x2))) def ifloordiv[T: Array](self, x1: T, x2: int | float | Array) -> T: - x1.value = tf.math.floordiv(x1.value, _unwrap(x2)) + x1.value = tf.math.floordiv(x1.value, unwrap(x2)) return x1 def remainder(self, x1: int | float | Array, x2: int | float | Array) -> Array: - return Array(tf.math.floormod(_unwrap(x1), _unwrap(x2))) + return Array(tf.math.floormod(unwrap(x1), unwrap(x2))) def imod[T: Array](self, x1: T, x2: int | float | Array) -> T: - x1.value = tf.math.floormod(x1.value, _unwrap(x2)) + x1.value = tf.math.floormod(x1.value, unwrap(x2)) return x1 def pow(self, x1: int | float | complex | Array, x2: int | float | complex | Array) -> Array: - return Array(tf.pow(_unwrap(x1), _unwrap(x2))) + return Array(tf.pow(unwrap(x1), unwrap(x2))) def ipow[T: Array](self, x1: T, x2: int | float | complex | Array) -> T: - x1.value = tf.pow(x1.value, _unwrap(x2)) + x1.value = tf.pow(x1.value, unwrap(x2)) return x1 def negative(self, x: Array) -> Array: @@ -270,22 +266,22 @@ def sqrt(self, x: Array) -> Array: # Comparisons def equal(self, x1: int | float | complex | Array, x2: int | float | complex | Array) -> Array: - return Array(tf.equal(_unwrap(x1), _unwrap(x2))) + return Array(tf.equal(unwrap(x1), unwrap(x2))) def not_equal(self, x1: int | float | complex | Array, x2: int | float | complex | Array) -> Array: - return Array(tf.not_equal(_unwrap(x1), _unwrap(x2))) + return Array(tf.not_equal(unwrap(x1), unwrap(x2))) def less(self, x1: int | float | complex | Array, x2: int | float | complex | Array) -> Array: - return Array(tf.less(_unwrap(x1), _unwrap(x2))) + return Array(tf.less(unwrap(x1), unwrap(x2))) def less_equal(self, x1: int | float | complex | Array, x2: int | float | complex | Array) -> Array: - return Array(tf.less_equal(_unwrap(x1), _unwrap(x2))) + return Array(tf.less_equal(unwrap(x1), unwrap(x2))) def greater(self, x1: int | float | complex | Array, x2: int | float | complex | Array) -> Array: - return Array(tf.greater(_unwrap(x1), _unwrap(x2))) + return Array(tf.greater(unwrap(x1), unwrap(x2))) def greater_equal(self, x1: int | float | complex | Array, x2: int | float | complex | Array) -> Array: - return Array(tf.greater_equal(_unwrap(x1), _unwrap(x2))) + return Array(tf.greater_equal(unwrap(x1), unwrap(x2))) # Bitwise — TF's native ``&`` dispatches to ``tf.math.logical_and`` for bool # tensors and ``tf.bitwise.bitwise_and`` for int tensors, matching numpy/torch/jax @@ -293,41 +289,41 @@ def greater_equal(self, x1: int | float | complex | Array, x2: int | float | com # us to one dtype family. def bitwise_and(self, x1: bool | int | Array, x2: bool | int | Array) -> Array: - return Array(_unwrap(x1) & _unwrap(x2)) + return Array(unwrap(x1) & unwrap(x2)) def iand[T: Array](self, x1: T, x2: bool | int | Array) -> T: - x1.value &= _unwrap(x2) + x1.value &= unwrap(x2) return x1 def bitwise_invert(self, x: Array) -> Array: return Array(~x.value) def bitwise_or(self, x1: bool | int | Array, x2: bool | int | Array) -> Array: - return Array(_unwrap(x1) | _unwrap(x2)) + return Array(unwrap(x1) | unwrap(x2)) def ior[T: Array](self, x1: T, x2: bool | int | Array) -> T: - x1.value |= _unwrap(x2) + x1.value |= unwrap(x2) return x1 def bitwise_xor(self, x1: bool | int | Array, x2: bool | int | Array) -> Array: - return Array(_unwrap(x1) ^ _unwrap(x2)) + return Array(unwrap(x1) ^ unwrap(x2)) def ixor[T: Array](self, x1: T, x2: bool | int | Array) -> T: - x1.value ^= _unwrap(x2) + x1.value ^= unwrap(x2) return x1 def bitwise_left_shift(self, x1: int | Array, x2: int | Array) -> Array: - return Array(tf.bitwise.left_shift(_unwrap(x1), _unwrap(x2))) + return Array(tf.bitwise.left_shift(unwrap(x1), unwrap(x2))) def ilshift[T: Array](self, x1: T, x2: int | Array) -> T: - x1.value = tf.bitwise.left_shift(x1.value, _unwrap(x2)) + x1.value = tf.bitwise.left_shift(x1.value, unwrap(x2)) return x1 def bitwise_right_shift(self, x1: int | Array, x2: int | Array) -> Array: - return Array(tf.bitwise.right_shift(_unwrap(x1), _unwrap(x2))) + return Array(tf.bitwise.right_shift(unwrap(x1), unwrap(x2))) def irshift[T: Array](self, x1: T, x2: int | Array) -> T: - x1.value = tf.bitwise.right_shift(x1.value, _unwrap(x2)) + x1.value = tf.bitwise.right_shift(x1.value, unwrap(x2)) return x1 # Operators @@ -336,7 +332,7 @@ def sign(self, x: Array) -> Array: return Array(tf.sign(x.value)) def maximum(self, x1: int | float | complex | Array, x2: int | float | complex | Array) -> Array: - return Array(tf.maximum(_unwrap(x1), _unwrap(x2))) + return Array(tf.maximum(unwrap(x1), unwrap(x2))) def argmax(self, x: Array, axis: int | None = None, keepdims: bool = False) -> Array: v = x.value @@ -371,7 +367,7 @@ def set_item(self, x: Array, key: ArrayKey, value: bool | int | float | complex # that hammer set_item in tight loops should consider numpy or pytorch. original = x.value np_array = original.numpy().copy() - np_array[key] = np.asarray(_unwrap(value)) + np_array[key] = np.asarray(unwrap(value)) with tf.device(self._native_device): x.value = tf.convert_to_tensor(np_array, dtype=original.dtype) From 4066770979d923ed076f44178a87384c8d2b1db5 Mon Sep 17 00:00:00 2001 From: nicola-bastianello Date: Thu, 18 Jun 2026 20:06:26 +0200 Subject: [PATCH 14/23] enh: dtypes docstring --- decent_array/types/_dtypes.py | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) diff --git a/decent_array/types/_dtypes.py b/decent_array/types/_dtypes.py index 2c46c97..7feccd2 100644 --- a/decent_array/types/_dtypes.py +++ b/decent_array/types/_dtypes.py @@ -155,7 +155,19 @@ def __hash__(self) -> int: def dtypes(*, kind: str | tuple[str, ...] | None = None) -> dict[str, dtype]: - """Return a dictionary of available dtypes.""" + """ + Return a dictionary of available dtypes. + + Args: + kind: kind of dtypes to be returned, either one string or a tuple of strings; available kinds are: `bool`, + `signed integer`, `unsigned integer`, `integral`, `real floating`, `complex floating`, `numeric`. If kind is + None, all available dtypes are included. If `kind` does not match any of the supported strings, an empty + dictionary is returned. + + Returns: + A dictionary of available dtypes, keyed by name, and filtered by `kind`. + + """ if kind is None: dtypes = _ALL_DTYPES elif isinstance(kind, str): From 5f68671ffc898a73d702a1135e2c1c50f2ff09f9 Mon Sep 17 00:00:00 2001 From: nicola-bastianello Date: Fri, 19 Jun 2026 10:16:20 +0200 Subject: [PATCH 15/23] fix: remove ruff preview rules --- decent_array/_array.py | 4 ++-- decent_array/interoperability/_abstracts/backend.py | 2 +- decent_array/interoperability/_backend_manager.py | 2 +- decent_array/interoperability/_jax/jax_backend.py | 2 +- decent_array/interoperability/_numpy/numpy_backend.py | 2 +- decent_array/interoperability/_pytorch/pytorch_backend.py | 2 +- .../interoperability/_tensorflow/tensorflow_backend.py | 2 +- pyproject.toml | 2 +- 8 files changed, 9 insertions(+), 9 deletions(-) diff --git a/decent_array/_array.py b/decent_array/_array.py index ca530a0..cc73f1c 100644 --- a/decent_array/_array.py +++ b/decent_array/_array.py @@ -49,7 +49,7 @@ def _update_backend(backend: Backend | None) -> None: register_backend_listener(_update_backend) -class Array: # noqa: PLR0904 +class Array: """ Wrapper around a single backend-native array. @@ -402,7 +402,7 @@ def ndim(self) -> int: @property def dtype(self) -> dtype: - """Return dtype of the Array.""" # noqa: DOC501 + """Return dtype of the Array.""" dtype = _BACKEND_DTYPE_TO_DTYPE.get(self.value.dtype) if dtype is None: diff --git a/decent_array/interoperability/_abstracts/backend.py b/decent_array/interoperability/_abstracts/backend.py index 4160f74..1912a02 100644 --- a/decent_array/interoperability/_abstracts/backend.py +++ b/decent_array/interoperability/_abstracts/backend.py @@ -28,7 +28,7 @@ from decent_array.types._dtypes import dtype -class Backend(ABC): # noqa: PLR0904 +class Backend(ABC): """ Abstract base class for a backend. diff --git a/decent_array/interoperability/_backend_manager.py b/decent_array/interoperability/_backend_manager.py index 7c30169..c897fea 100644 --- a/decent_array/interoperability/_backend_manager.py +++ b/decent_array/interoperability/_backend_manager.py @@ -46,7 +46,7 @@ def set_backend( is already active in this context. ImportError: If the backend module cannot be imported (e.g. due to a missing optional dependency). - """ # noqa: DOC502 + """ requested = _normalize(backend) requested_device = device if isinstance(device, Devices) else Devices(device) diff --git a/decent_array/interoperability/_jax/jax_backend.py b/decent_array/interoperability/_jax/jax_backend.py index a84dd75..744979c 100644 --- a/decent_array/interoperability/_jax/jax_backend.py +++ b/decent_array/interoperability/_jax/jax_backend.py @@ -27,7 +27,7 @@ from decent_array.types._dtypes import _ALL_DTYPES, dtype -class JaxBackend(Backend): # noqa: PLR0904 +class JaxBackend(Backend): """JAX implementation of :class:`Backend`.""" def __init__(self, device: Devices = Devices.CPU) -> None: diff --git a/decent_array/interoperability/_numpy/numpy_backend.py b/decent_array/interoperability/_numpy/numpy_backend.py index d94727f..3fe66d0 100644 --- a/decent_array/interoperability/_numpy/numpy_backend.py +++ b/decent_array/interoperability/_numpy/numpy_backend.py @@ -22,7 +22,7 @@ from decent_array.types._dtypes import _ALL_DTYPES, dtype -class NumpyBackend(Backend): # noqa: PLR0904 +class NumpyBackend(Backend): """NumPy implementation of :class:`Backend`.""" def __init__(self, device: Devices = Devices.CPU) -> None: diff --git a/decent_array/interoperability/_pytorch/pytorch_backend.py b/decent_array/interoperability/_pytorch/pytorch_backend.py index 2936944..36e1789 100644 --- a/decent_array/interoperability/_pytorch/pytorch_backend.py +++ b/decent_array/interoperability/_pytorch/pytorch_backend.py @@ -22,7 +22,7 @@ from decent_array.types._dtypes import _ALL_DTYPES, dtype -class PyTorchBackend(Backend): # noqa: PLR0904 +class PyTorchBackend(Backend): """PyTorch implementation of :class:`Backend`.""" def __init__(self, device: Devices = Devices.CPU) -> None: diff --git a/decent_array/interoperability/_tensorflow/tensorflow_backend.py b/decent_array/interoperability/_tensorflow/tensorflow_backend.py index 498b9d9..fe5b150 100644 --- a/decent_array/interoperability/_tensorflow/tensorflow_backend.py +++ b/decent_array/interoperability/_tensorflow/tensorflow_backend.py @@ -25,7 +25,7 @@ from decent_array.types._dtypes import _ALL_DTYPES, dtype -class TensorflowBackend(Backend): # noqa: PLR0904 +class TensorflowBackend(Backend): """TensorFlow implementation of :class:`Backend`.""" def __init__(self, device: Devices = Devices.CPU) -> None: diff --git a/pyproject.toml b/pyproject.toml index 398c054..41daa86 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -194,5 +194,5 @@ lint.ignore = [ "ICN001", # unconventional-import-alias, complains about common import aliases like `import numpy as np` but it doesn't work for most libraries "PYI041", # Removes warning for redundant types in cases like int | float | complex where complex already covers int and float, but we want to keep the union for readability and explicitness about supported types. ] -preview = true +preview = false line-length = 120 From b67c251f5b112fae0287f6175a604ace4a25d34f Mon Sep 17 00:00:00 2001 From: nicola-bastianello Date: Wed, 1 Jul 2026 15:36:30 +0200 Subject: [PATCH 16/23] fix: address PR comments --- .../interoperability/_jax/jax_backend.py | 4 +-- .../interoperability/_numpy/numpy_backend.py | 4 +-- .../_pytorch/pytorch_backend.py | 21 ++++++++--- .../_tensorflow/tensorflow_backend.py | 4 +-- decent_array/types/_dtypes.py | 36 +++++++++---------- docs/source/user.rst | 4 +-- 6 files changed, 40 insertions(+), 33 deletions(-) diff --git a/decent_array/interoperability/_jax/jax_backend.py b/decent_array/interoperability/_jax/jax_backend.py index 744979c..74ca145 100644 --- a/decent_array/interoperability/_jax/jax_backend.py +++ b/decent_array/interoperability/_jax/jax_backend.py @@ -24,7 +24,7 @@ from decent_array.interoperability._abstracts import Backend from decent_array.interoperability._backend_manager import register_backend from decent_array.types import ArrayKey, ArrayTypes, Devices, Frameworks -from decent_array.types._dtypes import _ALL_DTYPES, dtype +from decent_array.types._dtypes import dtype class JaxBackend(Backend): @@ -128,7 +128,7 @@ def diagonal(self, x: Array, offset: int = 0) -> Array: return Array(jnp.diagonal(x.value, offset=offset)) def astype(self, x: Array, dtype: dtype) -> Array: - if dtype not in _ALL_DTYPES.values(): + if not dtype.available: raise ValueError(f"Unsupported dtype '{dtype}' for JAX backend.") return Array(jnp.asarray(x.value, dtype=dtype.backend_dtype)) diff --git a/decent_array/interoperability/_numpy/numpy_backend.py b/decent_array/interoperability/_numpy/numpy_backend.py index 3fe66d0..a49b644 100644 --- a/decent_array/interoperability/_numpy/numpy_backend.py +++ b/decent_array/interoperability/_numpy/numpy_backend.py @@ -19,7 +19,7 @@ from decent_array.interoperability._abstracts import Backend from decent_array.interoperability._backend_manager import register_backend from decent_array.types import ArrayKey, ArrayTypes, Devices, Frameworks -from decent_array.types._dtypes import _ALL_DTYPES, dtype +from decent_array.types._dtypes import dtype class NumpyBackend(Backend): @@ -123,7 +123,7 @@ def diagonal(self, x: Array, offset: int = 0) -> Array: return Array(np.diagonal(x.value, offset=offset)) def astype(self, x: Array, dtype: dtype) -> Array: - if dtype not in _ALL_DTYPES.values(): + if not dtype.available: raise ValueError(f"Unsupported dtype '{dtype}' for NumPy backend.") return Array(np.asarray(x.value, dtype=dtype.backend_dtype)) diff --git a/decent_array/interoperability/_pytorch/pytorch_backend.py b/decent_array/interoperability/_pytorch/pytorch_backend.py index 36e1789..688e9fd 100644 --- a/decent_array/interoperability/_pytorch/pytorch_backend.py +++ b/decent_array/interoperability/_pytorch/pytorch_backend.py @@ -19,7 +19,7 @@ from decent_array.interoperability._abstracts import Backend from decent_array.interoperability._backend_manager import register_backend from decent_array.types import ArrayKey, ArrayTypes, Devices, Frameworks -from decent_array.types._dtypes import _ALL_DTYPES, dtype +from decent_array.types._dtypes import dtype class PyTorchBackend(Backend): @@ -138,7 +138,7 @@ def diagonal(self, x: Array, offset: int = 0) -> Array: return Array(torch.diagonal(x.value, offset=offset)) def astype(self, x: Array, dtype: dtype) -> Array: - if dtype not in _ALL_DTYPES.values(): + if not dtype.available: raise ValueError(f"Unsupported dtype '{dtype}' for PyTorch backend.") return Array(x.value.to(dtype=dtype.backend_dtype)) @@ -437,7 +437,9 @@ def float32(self) -> torch.dtype: return torch.float32 @property - def float64(self) -> torch.dtype: + def float64(self) -> Any: # noqa: ANN401 + if self.device == Devices.MPS: + return None return torch.float64 @property @@ -461,8 +463,17 @@ def quint8(self) -> Any: # noqa: ANN401 return torch.quint8 @property - def bfloat16(self) -> torch.dtype: - return torch.bfloat16 + def bfloat16(self) -> Any: # noqa: ANN401 + if self.device == Devices.GPU: + return torch.bfloat16 if torch.cuda.is_bf16_supported() else None + if self.device == Devices.MPS: + return None + capabilities = torch.cpu.get_capabilities() + return ( + torch.bfloat16 + if any(tag in capabilities for tag in ("bf16", "avx512_bf16", "amx_bf16", "sve_bf16")) + else None + ) # Constants diff --git a/decent_array/interoperability/_tensorflow/tensorflow_backend.py b/decent_array/interoperability/_tensorflow/tensorflow_backend.py index fe5b150..7acc1a9 100644 --- a/decent_array/interoperability/_tensorflow/tensorflow_backend.py +++ b/decent_array/interoperability/_tensorflow/tensorflow_backend.py @@ -22,7 +22,7 @@ from decent_array.interoperability._abstracts import Backend from decent_array.interoperability._backend_manager import register_backend from decent_array.types import ArrayKey, ArrayTypes, Devices, Frameworks -from decent_array.types._dtypes import _ALL_DTYPES, dtype +from decent_array.types._dtypes import dtype class TensorflowBackend(Backend): @@ -143,7 +143,7 @@ def diagonal(self, x: Array, offset: int = 0) -> Array: return Array(tf.linalg.diag_part(v, k=offset)) def astype(self, x: Array, dtype: dtype) -> Array: - if dtype not in _ALL_DTYPES.values(): + if not dtype.available: raise ValueError(f"Unsupported dtype '{dtype}' for TensorFlow backend.") return Array(tf.cast(x.value, dtype=dtype.backend_dtype)) diff --git a/decent_array/types/_dtypes.py b/decent_array/types/_dtypes.py index 7feccd2..f8dec4c 100644 --- a/decent_array/types/_dtypes.py +++ b/decent_array/types/_dtypes.py @@ -83,52 +83,46 @@ def __hash__(self) -> int: # instantiate all the supported dtypes bool_ = dtype("bool_") -_BOOL_DTYPES = {"bool_": bool_} +_BOOL_DTYPES = {bool_} int8 = dtype("int8") int16 = dtype("int16") int32 = dtype("int32") int64 = dtype("int64") -_SIGNED_INT_DTYPES = {"int8": int8, "int16": int16, "int32": int32, "int64": int64} +_SIGNED_INT_DTYPES = {int8, int16, int32, int64} uint8 = dtype("uint8") uint16 = dtype("uint16") uint32 = dtype("uint32") uint64 = dtype("uint64") -_UNSIGNED_INT_DTYPES = {"uint8": uint8, "uint16": uint16, "uint32": uint32, "uint64": uint64} +_UNSIGNED_INT_DTYPES = {uint8, uint16, uint32, uint64} float16 = dtype("float16") bfloat16 = dtype("bfloat16") float32 = dtype("float32") float64 = dtype("float64") float128 = dtype("float128") -_REAL_FLOATING_DTYPES = { - "float16": float16, - "bfloat16": bfloat16, - "float32": float32, - "float64": float64, - "float128": float128, -} +_REAL_FLOATING_DTYPES = {float16, bfloat16, float32, float64, float128} complex64 = dtype("complex64") complex128 = dtype("complex128") complex256 = dtype("complex256") -_COMPLEX_FLOATING_DTYPES = {"complex64": complex64, "complex128": complex128, "complex256": complex256} +_COMPLEX_FLOATING_DTYPES = {complex64, complex128, complex256} qint8 = dtype("qint8") qint16 = dtype("qint16") qint32 = dtype("qint32") -_QUANTIZED_SIGNED_INT_DTYPES = {"qint8": qint8, "qint16": qint16, "qint32": qint32} +_QUANTIZED_SIGNED_INT_DTYPES = {qint8, qint16, qint32} quint8 = dtype("quint8") quint16 = dtype("quint16") -_QUANTIZED_UNSIGNED_INT_DTYPES = {"quint8": quint8, "quint16": quint16} +_QUANTIZED_UNSIGNED_INT_DTYPES = {quint8, quint16} unicode_ = dtype("unicode_") bytes_ = dtype("bytes_") object_ = dtype("object_") void = dtype("void") -_MISCELLANEOUS_DTYPES = {"unicode_": unicode_, "bytes_": bytes_, "object_": object_, "void": void} +_MISCELLANEOUS_DTYPES = {unicode_, bytes_, object_, void} _INTEGRAL_DTYPES = ( _SIGNED_INT_DTYPES | _UNSIGNED_INT_DTYPES | _QUANTIZED_SIGNED_INT_DTYPES | _QUANTIZED_UNSIGNED_INT_DTYPES @@ -137,10 +131,12 @@ def __hash__(self) -> int: _ALL_DTYPES = _BOOL_DTYPES | _NUMERIC_DTYPES | _MISCELLANEOUS_DTYPES +_AVAILABLE_DTYPES = {dt for dt in _ALL_DTYPES if dt.available} + + _BACKEND_DTYPE_TO_DTYPE: dict[Any, dtype] = {} -for dt in _ALL_DTYPES.values(): - if dt.available: - _BACKEND_DTYPE_TO_DTYPE[dt._backend_dtype] = dt # noqa: SLF001 +for dt in _AVAILABLE_DTYPES: + _BACKEND_DTYPE_TO_DTYPE[dt._backend_dtype] = dt # noqa: SLF001 _ALIASES = { @@ -171,8 +167,8 @@ def dtypes(*, kind: str | tuple[str, ...] | None = None) -> dict[str, dtype]: if kind is None: dtypes = _ALL_DTYPES elif isinstance(kind, str): - dtypes = _ALIASES.get(kind, {}) + dtypes = _ALIASES.get(kind, set()) else: - dtypes = {k: v for alias in kind if alias in _ALIASES for k, v in _ALIASES[alias].items()} + dtypes = {dt for alias in kind if alias in _ALIASES for dt in _ALIASES[alias]} - return {name: dt for name, dt in dtypes.items() if dt.available} + return {dt.name: dt for dt in dtypes if dt.available} diff --git a/docs/source/user.rst b/docs/source/user.rst index 235744e..f31d0b9 100644 --- a/docs/source/user.rst +++ b/docs/source/user.rst @@ -182,7 +182,7 @@ reliable way is to try operations that involve the dtype and observe if the fram - ✓ - ✓ - ✓ - - + - PyTorch support is limited/experimental * - ``float32`` - ✓ - ✓ @@ -194,7 +194,7 @@ reliable way is to try operations that involve the dtype and observe if the fram - ⚠️ - ✓ - ✓ - - JAX requires ``jax_enable_x64=True`` + - JAX requires ``jax_enable_x64=True``; PyTorch does not support on MPS * - ``float128`` - ⚠️ - ✗ From 3d1f1f19c4a4742b04085c9e661244ee13e60e1f Mon Sep 17 00:00:00 2001 From: nicola-bastianello Date: Wed, 1 Jul 2026 15:49:11 +0200 Subject: [PATCH 17/23] fix: added version control --- .github/dependabot.yml | 13 +++++++++++++ pyproject.toml | 30 +++++++++++++++--------------- 2 files changed, 28 insertions(+), 15 deletions(-) create mode 100644 .github/dependabot.yml diff --git a/.github/dependabot.yml b/.github/dependabot.yml new file mode 100644 index 0000000..57e8244 --- /dev/null +++ b/.github/dependabot.yml @@ -0,0 +1,13 @@ +version: 2 +updates: + # version update of dependencies (as specified in pyproject.toml) + - package-ecosystem: "pip" + directory: "/" + schedule: + interval: "weekly" + + # version update of github action tools (as specified in workflows/ci.yaml) + - package-ecosystem: "github-actions" + directory: "/" + schedule: + interval: "weekly" \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index 41daa86..4c6d853 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -22,31 +22,31 @@ Issues = "https://github.com/team-decent/decent-array/issues" [project.optional-dependencies] dev = [ - "hatch", - "mypy", - "pytest", - "ruff", - "torch", - "torchvision", - "types-tensorflow", + "hatch>=1.14,<2", + "mypy==2.1.0", + "pytest>=9,<10", + "ruff==0.15.20", + "torch>=2.5,<3", + "torchvision>=0.20,<1", + "types-tensorflow>=2.18,<3", ] dev-cpu = [ - "tensorflow", - "jax", + "tensorflow>=2.18,<3", + "jax>=0.4,<1", ] dev-gpu = [ - "tensorflow[and-cuda]", - "jax[cuda12]", + "tensorflow[and-cuda]>=2.18,<3", + "jax[cuda12]>=0.4,<1", ] sphinx-tools = [ - "sphinx", - "pydata-sphinx-theme", - "sphinxcontrib-bibtex", + "sphinx==9.1.0", + "pydata-sphinx-theme>=0.16,<1", + "sphinxcontrib-bibtex>=2.6,<3", ] [build-system] -requires = ["hatchling", "hatch-mypyc"] +requires = ["hatchling>=1.26,<2", "hatch-mypyc>=0.16.0,<1"] build-backend = "hatchling.build" # Compile hot-path modules with mypyc when building wheels. End users installing from # a wheel get a precompiled .so; ``pip install -e .`` and source builds also run this From e259113153b94f4af0ed2367d835b5408dbc8de7 Mon Sep 17 00:00:00 2001 From: nicola-bastianello Date: Mon, 6 Jul 2026 14:30:06 +0200 Subject: [PATCH 18/23] fix: properly lock versions, improve pyproject.toml, fix numpy-related mypy issues --- .../interoperability/_numpy/numpy_backend.py | 46 +++++++--- decent_array/types/_types.py | 8 +- docs/requirements.txt | 3 + docs/sphinx_theme.txt | 2 - pyproject.toml | 84 +++++++++++-------- readthedocs.yaml | 2 +- 6 files changed, 95 insertions(+), 50 deletions(-) create mode 100644 docs/requirements.txt delete mode 100644 docs/sphinx_theme.txt diff --git a/decent_array/interoperability/_numpy/numpy_backend.py b/decent_array/interoperability/_numpy/numpy_backend.py index a49b644..1cabeaa 100644 --- a/decent_array/interoperability/_numpy/numpy_backend.py +++ b/decent_array/interoperability/_numpy/numpy_backend.py @@ -9,7 +9,7 @@ from collections.abc import Sequence from copy import deepcopy -from typing import Any +from typing import Any, cast import numpy as np from numpy.typing import NDArray @@ -146,27 +146,47 @@ def vector_norm( keepdims: bool = False, ord: int | float = 2, # noqa: A002 ) -> Array: - return Array(np.linalg.norm(x.value, ord=ord, axis=axis, keepdims=keepdims)) + # axis in np.linalg.vector_norm seems to allow for int | tuple[int, ...] | None at runtime, + # but is still typed as int | tuple[int, int] | None, hence the ignore + return Array(np.linalg.vector_norm(x.value, ord=ord, axis=axis, keepdims=keepdims)) # type: ignore[arg-type] # Math reductions def sum(self, x: Array, axis: int | tuple[int, ...] | None = None, keepdims: bool = False) -> Array: - return Array(np.sum(x.value, axis=axis, keepdims=keepdims)) + v = cast("np.ndarray[Any, Any]", x.value) + if keepdims: + return Array(np.sum(v, axis=axis, keepdims=True)) + return Array(np.sum(v, axis=axis)) def mean(self, x: Array, axis: int | tuple[int, ...] | None = None, keepdims: bool = False) -> Array: - return Array(np.mean(x.value, axis=axis, keepdims=keepdims)) + v = cast("np.ndarray[Any, Any]", x.value) + if keepdims: + return Array(np.mean(v, axis=axis, keepdims=True)) + return Array(np.mean(v, axis=axis)) def min(self, x: Array, axis: int | tuple[int, ...] | None = None, keepdims: bool = False) -> Array: - return Array(np.min(x.value, axis=axis, keepdims=keepdims)) + v = cast("np.ndarray[Any, Any]", x.value) + if keepdims: + return Array(np.min(v, axis=axis, keepdims=True)) + return Array(np.min(v, axis=axis)) def max(self, x: Array, axis: int | tuple[int, ...] | None = None, keepdims: bool = False) -> Array: - return Array(np.max(x.value, axis=axis, keepdims=keepdims)) + v = cast("np.ndarray[Any, Any]", x.value) + if keepdims: + return Array(np.max(v, axis=axis, keepdims=True)) + return Array(np.max(v, axis=axis)) def any(self, x: Array, axis: int | tuple[int, ...] | None = None, keepdims: bool = False) -> bool: - return bool(np.any(x.value, axis=axis, keepdims=keepdims)) + v = cast("np.ndarray[Any, Any]", x.value) + if keepdims: + return bool(np.any(v, axis=axis, keepdims=True)) + return bool(np.any(v, axis=axis)) def all(self, x: Array, axis: int | tuple[int, ...] | None = None, keepdims: bool = False) -> bool: - return bool(np.all(x.value, axis=axis, keepdims=keepdims)) + v = cast("np.ndarray[Any, Any]", x.value) + if keepdims: + return bool(np.all(v, axis=axis, keepdims=True)) + return bool(np.all(v, axis=axis)) # Math elementwise — operands may be Array or scalar (operator dunders pass either). # ``Array | float`` covers both: PEP 484's numeric tower implicitly admits ``int``. @@ -298,10 +318,16 @@ def maximum(self, x1: int | float | complex | Array, x2: int | float | complex | return Array(np.maximum(unwrap(x1), unwrap(x2))) def argmax(self, x: Array, axis: int | None = None, keepdims: bool = False) -> Array: - return Array(np.argmax(x.value, axis=axis, keepdims=keepdims)) + v = cast("np.ndarray[Any, Any]", x.value) + if keepdims: + return Array(np.argmax(v, axis=axis, keepdims=True)) + return Array(np.argmax(v, axis=axis)) def argmin(self, x: Array, axis: int | None = None, keepdims: bool = False) -> Array: - return Array(np.argmin(x.value, axis=axis, keepdims=keepdims)) + v = cast("np.ndarray[Any, Any]", x.value) + if keepdims: + return Array(np.argmin(v, axis=axis, keepdims=True)) + return Array(np.argmin(v, axis=axis)) def set_item(self, x: Array, key: ArrayKey, value: bool | int | float | complex | Array) -> None: x.value[key] = unwrap(value) diff --git a/decent_array/types/_types.py b/decent_array/types/_types.py index ff5a081..10c6a1b 100644 --- a/decent_array/types/_types.py +++ b/decent_array/types/_types.py @@ -5,9 +5,10 @@ from enum import Enum from typing import TYPE_CHECKING, SupportsIndex, TypeAlias, Union +import numpy + if TYPE_CHECKING: import jax - import numpy import tensorflow as tf import torch @@ -20,10 +21,9 @@ PyTorch tensors, TensorFlow tensors, and JAX arrays. """ -ArrayTypes: TypeAlias = bool | int | float | complex | ArrayLike # noqa: UP040 +ArrayTypes: TypeAlias = bool | int | float | complex | numpy.generic | ArrayLike # noqa: UP040 """ -Type alias for supported types for optimization variables in decent-array, -including array-like types and scalars. +Type alias for supported scalar/array types in decent-array. """ ArrayKey: TypeAlias = ( # noqa: UP040 diff --git a/docs/requirements.txt b/docs/requirements.txt new file mode 100644 index 0000000..8e24718 --- /dev/null +++ b/docs/requirements.txt @@ -0,0 +1,3 @@ +sphinx==9.1.0 +pydata-sphinx-theme==0.19.0 +sphinxcontrib-bibtex==2.7.0 \ No newline at end of file diff --git a/docs/sphinx_theme.txt b/docs/sphinx_theme.txt deleted file mode 100644 index 6246669..0000000 --- a/docs/sphinx_theme.txt +++ /dev/null @@ -1,2 +0,0 @@ -pydata-sphinx-theme -sphinxcontrib-bibtex diff --git a/pyproject.toml b/pyproject.toml index 4c6d853..609a474 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -12,7 +12,7 @@ classifiers = [ ] license = "AGPL-3.0-only" dependencies = [ - "numpy", + "numpy>=2,<3", ] [project.urls] @@ -21,27 +21,36 @@ Source = "https://github.com/team-decent/decent-array" Issues = "https://github.com/team-decent/decent-array/issues" [project.optional-dependencies] -dev = [ - "hatch>=1.14,<2", - "mypy==2.1.0", - "pytest>=9,<10", - "ruff==0.15.20", +torch = [ "torch>=2.5,<3", "torchvision>=0.20,<1", - "types-tensorflow>=2.18,<3", ] -dev-cpu = [ +tf = [ "tensorflow>=2.18,<3", - "jax>=0.4,<1", ] -dev-gpu = [ +tf-gpu = [ "tensorflow[and-cuda]>=2.18,<3", +] +jax = [ + "jax>=0.4,<1", +] +jax-gpu = [ "jax[cuda12]>=0.4,<1", ] -sphinx-tools = [ - "sphinx==9.1.0", - "pydata-sphinx-theme>=0.16,<1", - "sphinxcontrib-bibtex>=2.6,<3", + +[dependency-groups] +lint = ["ruff==0.15.20"] +typecheck = [ + "mypy==2.1.0", + "types-tensorflow>=2.18,<3", + "microsoft-python-type-stubs @ git+https://github.com/microsoft/python-type-stubs.git@main", +] +test = ["pytest>=9,<10"] +dev = [ + "hatch>=1.14,<2", + { include-group = "lint" }, + { include-group = "typecheck" }, + { include-group = "test" }, ] @@ -55,7 +64,7 @@ build-backend = "hatchling.build" # the ``.py`` source (dev, dev-gpu, sphinx, mypy, ruff) opt out via # ``HATCH_BUILD_NO_HOOKS=true`` to keep their installs fast. [tool.hatch.build.targets.wheel.hooks.mypyc] -dependencies = ["hatch-mypyc"] +dependencies = ["hatch-mypyc>=0.16.0,<1"] # Compile the Array wrapper plus the entire iop2 package as one mypyc group: keeping # them in a single compilation unit means cross-module calls (Array → backend, iop # function → backend) are direct compiled-to-compiled calls. New framework backends @@ -74,31 +83,26 @@ package = "editable" [tool.tox.env.dev] description = "Generate dev venv with all dependencies, active with `source .tox/dev/bin/activate`" set_env = { HATCH_BUILD_NO_HOOKS = "true" } -deps = [ - ".[dev]", - ".[dev-cpu]", - ".[sphinx-tools]", - "git+https://github.com/microsoft/python-type-stubs.git@main" +deps = ["-rdocs/requirements.txt"] +commands_pre = [ + ["python", "-m", "pip", "install", "--group", "dev"], + ["python", "-m", "pip", "install", ".[torch,tf,jax]"], ] -package = "editable" [tool.tox.env.dev-gpu] description = "Generate dev venv with all dependencies including GPU support, active with `source .tox/dev-gpu/bin/activate`" set_env = { PIP_EXTRA_INDEX_URL = "https://download.pytorch.org/whl/cu128", HATCH_BUILD_NO_HOOKS = "true" } # torch 2.11 uses CUDA 13.0 which seems to have broken linalg.norm, use cuda 12.8 for now -deps = [ - ".[dev]", - ".[dev-gpu]", - ".[sphinx-tools]", - "git+https://github.com/microsoft/python-type-stubs.git@main" +deps = ["-rdocs/requirements.txt"] +commands_pre = [ + ["python", "-m", "pip", "install", "--group", "dev"], + ["python", "-m", "pip", "install", ".[torch,tf-gpu,jax-gpu]"], ] -package = "editable" [tool.tox.env.sphinx] description = "Generate rst and html files using sphinx" set_env = { HATCH_BUILD_NO_HOOKS = "true" } -skip_install = true allowlist_externals = ["tox"] -deps = [".[sphinx-tools]"] +deps = ["-rdocs/requirements.txt"] commands = [ ["tox", "-e", "clean-mypyc"], # ensure .py source files are used for doc generation, not mypyc-compiled .so ["sphinx-apidoc", "-o", "docs/source/api", "decent_array", "--separate", "--no-toc", "--templatedir=docs/source/_templates"], @@ -108,7 +112,11 @@ commands = [ [tool.tox.env.mypy] description = "Run mypy (static type checker)" set_env = { HATCH_BUILD_NO_HOOKS = "true" } -deps = [".[dev]", ".[dev-cpu]", "git+https://github.com/microsoft/python-type-stubs.git@main"] +deps = [] +commands_pre = [ + ["python", "-m", "pip", "install", "--group", "typecheck"], + ["python", "-m", "pip", "install", ".[torch,tf,jax]"], +] commands = [["mypy", "decent_array"]] [tool.tox.env.clean-mypyc] @@ -122,7 +130,10 @@ commands = [ [tool.tox.env.mypyc] description = "Compile hot-path modules in place with mypyc (drops .so files alongside the .py source)" -deps = ["mypy", "numpy", "setuptools"] +deps = ["setuptools>=81,<82"] +commands_pre = [ + ["python", "-m", "pip", "install", ".", "--group", "typecheck"], +] skip_install = true allowlist_externals = ["tox"] # First command clears any stale shared-runtime ``.so`` at the project root from a @@ -140,12 +151,19 @@ commands = [ [tool.tox.env.pytest] description = "Run pytest (test executor)" -deps = [".[dev]", ".[dev-cpu]"] +deps = [] +commands_pre = [ + ["python", "-m", "pip", "install", "--group", "test"], + ["python", "-m", "pip", "install", ".[torch,tf,jax]"], +] commands = [["pytest", "-rs"]] [tool.tox.env.ruff] description = "Run ruff (format and style checker)" -deps = ["ruff"] +deps = [] +commands_pre = [ + ["python", "-m", "pip", "install", "--group", "lint"], +] set_env = { HATCH_BUILD_NO_HOOKS = "true" } skip_install = true commands = [ diff --git a/readthedocs.yaml b/readthedocs.yaml index 3b75931..0f0534d 100644 --- a/readthedocs.yaml +++ b/readthedocs.yaml @@ -22,6 +22,6 @@ sphinx: python: install: - - requirements: docs/sphinx_theme.txt + - requirements: docs/requirements.txt - method: pip path: . From 2417bc78c3045a23167de5d5fb60ebce233e9da0 Mon Sep 17 00:00:00 2001 From: nicola-bastianello Date: Mon, 6 Jul 2026 16:00:04 +0200 Subject: [PATCH 19/23] enh: add item to array, base coercion dunders on item --- decent_array/_array.py | 19 ++++++--- decent_array/_utils.py | 5 +++ .../interoperability/_abstracts/backend.py | 8 ++++ .../interoperability/_jax/jax_backend.py | 14 ++++++- .../interoperability/_numpy/numpy_backend.py | 14 ++++++- .../_pytorch/pytorch_backend.py | 14 ++++++- .../_tensorflow/tensorflow_backend.py | 14 ++++++- tests/test_array.py | 42 ++++++++++++++++++- 8 files changed, 119 insertions(+), 11 deletions(-) diff --git a/decent_array/_array.py b/decent_array/_array.py index cc73f1c..861c16f 100644 --- a/decent_array/_array.py +++ b/decent_array/_array.py @@ -355,23 +355,23 @@ def __len__(self) -> int: def __float__(self) -> float: """Coerce a scalar array to a Python float.""" - return float(self._backend.squeeze(self).value) + return float(self.item()) def __bool__(self) -> bool: """Coerce a scalar array to a Python bool.""" - return bool(self._backend.squeeze(self).value) + return bool(self.item()) def __int__(self) -> int: """Coerce a scalar array to a Python int.""" - return int(self._backend.squeeze(self).value) + return int(self.item()) def __complex__(self) -> complex: """Coerce a scalar array to a Python complex.""" - return complex(self._backend.squeeze(self).value) + return complex(self.item()) def __index__(self) -> int: """Coerce a scalar array to a Python int.""" - return int(self._backend.squeeze(self).value) + return int(self.item()) # Repr ----------------------------------------------------------------- @@ -444,3 +444,12 @@ def device(self) -> Devices: def numpy(self) -> NDArray[Any]: """Return a NumPy array view of the array's data.""" return self._backend.to_numpy(self) + + def item(self) -> Any: # noqa: ANN401 + """ + Convert 0-dim array to Python scalar. + + Raises: + TypeError: if ``x`` is not 0-dimensional. + """ # noqa: DOC502 + return self._backend.to_scalar(self) diff --git a/decent_array/_utils.py b/decent_array/_utils.py index 559e3d5..adb0462 100644 --- a/decent_array/_utils.py +++ b/decent_array/_utils.py @@ -12,3 +12,8 @@ def unwrap(x: Any) -> Any: # noqa: ANN401 site without runtime benefit. """ return x.value if type(x) is Array else x + + +def is_scalar(x: Array) -> bool: + """Return True if ``x`` is a 0-dim Array.""" + return x.ndim == 0 diff --git a/decent_array/interoperability/_abstracts/backend.py b/decent_array/interoperability/_abstracts/backend.py index 1912a02..80e084c 100644 --- a/decent_array/interoperability/_abstracts/backend.py +++ b/decent_array/interoperability/_abstracts/backend.py @@ -92,6 +92,14 @@ def from_numpy_like(self, x: NDArray[Any], like: Array) -> Array: def asarray(self, x: bool | int | float | complex) -> Array: """Convert a Python scalar to an :class:`Array` on this backend.""" + @abstractmethod + def to_scalar(self, x: Array) -> Any: # noqa: ANN401 + """ + Convert a 0-dim array to a scalar. + + This method must use ``is_scalar(x)`` to raise when ``x`` is not 0-dim. + """ + @abstractmethod def stack(self, arrays: Sequence[Array], axis: int = 0) -> Array: """Stack a sequence of arrays along a new dimension.""" diff --git a/decent_array/interoperability/_jax/jax_backend.py b/decent_array/interoperability/_jax/jax_backend.py index 74ca145..3859b6c 100644 --- a/decent_array/interoperability/_jax/jax_backend.py +++ b/decent_array/interoperability/_jax/jax_backend.py @@ -20,7 +20,7 @@ from numpy.typing import NDArray from decent_array import Array -from decent_array._utils import unwrap +from decent_array._utils import is_scalar, unwrap from decent_array.interoperability._abstracts import Backend from decent_array.interoperability._backend_manager import register_backend from decent_array.types import ArrayKey, ArrayTypes, Devices, Frameworks @@ -85,6 +85,18 @@ def from_numpy_like(self, x: NDArray[Any], like: Array) -> Array: def asarray(self, x: bool | int | float | complex) -> Array: return Array(jnp.array(x, device=self._native_device)) + def to_scalar(self, x: Array) -> Any: # noqa: ANN401 + """ + Convert a 0-dim array to a scalar. + + Raises: + TypeError: if ``x`` is not 0-dimensional. + + """ + if not is_scalar(x): + raise TypeError("Only 0-dim arrays can be converted to Python scalars.") + return x.value.item() + def stack(self, arrays: Sequence[Array], axis: int = 0) -> Array: if len(arrays) == 0: raise ValueError("Cannot stack an empty sequence of arrays.") diff --git a/decent_array/interoperability/_numpy/numpy_backend.py b/decent_array/interoperability/_numpy/numpy_backend.py index 1cabeaa..d032ef4 100644 --- a/decent_array/interoperability/_numpy/numpy_backend.py +++ b/decent_array/interoperability/_numpy/numpy_backend.py @@ -15,7 +15,7 @@ from numpy.typing import NDArray from decent_array import Array -from decent_array._utils import unwrap +from decent_array._utils import is_scalar, unwrap from decent_array.interoperability._abstracts import Backend from decent_array.interoperability._backend_manager import register_backend from decent_array.types import ArrayKey, ArrayTypes, Devices, Frameworks @@ -80,6 +80,18 @@ def from_numpy_like(self, x: NDArray[Any], like: Array) -> Array: def asarray(self, x: bool | int | float | complex) -> Array: return Array(np.array(x)) + def to_scalar(self, x: Array) -> Any: # noqa: ANN401 + """ + Convert a 0-dim array to a scalar. + + Raises: + TypeError: if ``x`` is not 0-dimensional. + + """ # noqa: DOC501, DOC502 + if not is_scalar(x): + raise TypeError("Only 0-dim arrays can be converted to Python scalars.") + return x.value.item() + def stack(self, arrays: Sequence[Array], axis: int = 0) -> Array: if len(arrays) == 0: raise ValueError("Cannot stack an empty sequence of arrays.") diff --git a/decent_array/interoperability/_pytorch/pytorch_backend.py b/decent_array/interoperability/_pytorch/pytorch_backend.py index 688e9fd..7d39174 100644 --- a/decent_array/interoperability/_pytorch/pytorch_backend.py +++ b/decent_array/interoperability/_pytorch/pytorch_backend.py @@ -15,7 +15,7 @@ from numpy.typing import NDArray from decent_array import Array -from decent_array._utils import unwrap +from decent_array._utils import is_scalar, unwrap from decent_array.interoperability._abstracts import Backend from decent_array.interoperability._backend_manager import register_backend from decent_array.types import ArrayKey, ArrayTypes, Devices, Frameworks @@ -90,6 +90,18 @@ def from_numpy_like(self, x: NDArray[Any], like: Array) -> Array: def asarray(self, x: bool | int | float | complex) -> Array: return Array(torch.tensor(x, device=self._native_device)) + def to_scalar(self, x: Array) -> Any: # noqa: ANN401 + """ + Convert a 0-dim array to a scalar. + + Raises: + TypeError: if ``x`` is not 0-dimensional. + + """ + if not is_scalar(x): + raise TypeError("Only 0-dim arrays can be converted to Python scalars.") + return x.value.item() + def stack(self, arrays: Sequence[Array], axis: int = 0) -> Array: if len(arrays) == 0: raise ValueError("Cannot stack an empty sequence of arrays.") diff --git a/decent_array/interoperability/_tensorflow/tensorflow_backend.py b/decent_array/interoperability/_tensorflow/tensorflow_backend.py index 7acc1a9..438d6f1 100644 --- a/decent_array/interoperability/_tensorflow/tensorflow_backend.py +++ b/decent_array/interoperability/_tensorflow/tensorflow_backend.py @@ -18,7 +18,7 @@ from numpy.typing import NDArray from decent_array import Array -from decent_array._utils import unwrap +from decent_array._utils import is_scalar, unwrap from decent_array.interoperability._abstracts import Backend from decent_array.interoperability._backend_manager import register_backend from decent_array.types import ArrayKey, ArrayTypes, Devices, Frameworks @@ -95,6 +95,18 @@ def asarray(self, x: bool | int | float | complex) -> Array: # Its not a tf tensor but mypyc doesn't import tf so it complains about unsude type-ignores return Array(tf.convert_to_tensor(cast("tf.Tensor", x))) + def to_scalar(self, x: Array) -> Any: # noqa: ANN401 + """ + Convert a 0-dim array to a scalar. + + Raises: + TypeError: if ``x`` is not 0-dimensional. + + """ + if not is_scalar(x): + raise TypeError("Only 0-dim arrays can be converted to Python scalars.") + return x.value.numpy().item() + def stack(self, arrays: Sequence[Array], axis: int = 0) -> Array: if len(arrays) == 0: raise ValueError("Cannot stack an empty sequence of arrays.") diff --git a/tests/test_array.py b/tests/test_array.py index f2a6cb0..9ca8d3d 100644 --- a/tests/test_array.py +++ b/tests/test_array.py @@ -432,11 +432,34 @@ def test_len(backend: tuple) -> None: assert len(a) == 4 -def test_float_coercion(backend: tuple) -> None: - a = _create_array([2.5]) +def test_float_coercion_float(backend: tuple) -> None: + a = _create_array(2.5) assert float(a) == pytest.approx(2.5) +def test_float_coercion_bool(backend: tuple) -> None: + a = _create_array(1) + assert bool(a) == True + + a = _create_array(0) + assert bool(a) == False + + +def test_float_coercion_int(backend: tuple) -> None: + a = _create_array(10) + assert int(a) == 10 + + +def test_float_coercion_index(backend: tuple) -> None: + a = _create_array(10) + assert int(a) == 10 + + +def test_float_coercion_complex(backend: tuple) -> None: + a = iop.from_numpy(np.array(1 + 2 * 1j, dtype=np.complex64)) + assert complex(a) == 1 + 2 * 1j + + def test_repr(backend: tuple) -> None: a = _create_array([1.0, 2.0]) text = repr(a) @@ -529,3 +552,18 @@ def test_device_property(backend: tuple) -> None: _framework, device = backend a = iop.zeros((3,)) assert a.device == device + + +def test_item_0_dim(backend: tuple) -> None: + a = _create_array(1.0) + assert a.item() == 1.0 + + +def test_item_n_dim(backend: tuple) -> None: + a = _create_array([1.0]) + with pytest.raises(TypeError, match=r"Only 0-dim arrays"): + _ = a.item() + + a = _create_array([1.0, 2.0]) + with pytest.raises(TypeError, match=r"Only 0-dim arrays"): + _ = a.item() From 8542728dd0162af505715042b39e2c72c12fef49 Mon Sep 17 00:00:00 2001 From: nicola-bastianello Date: Mon, 6 Jul 2026 17:10:50 +0200 Subject: [PATCH 20/23] enh: unify error messaging --- decent_array/_array.py | 3 +- decent_array/_errors.py | 37 ++++++++++++++++++ .../interoperability/_iop/bit_operators.py | 13 ++++--- .../interoperability/_iop/comparison.py | 14 +++---- .../interoperability/_iop/creation.py | 12 +++--- decent_array/interoperability/_iop/linalg.py | 12 +++--- .../interoperability/_iop/manipulations.py | 38 +++++++++---------- decent_array/interoperability/_iop/math.py | 34 ++++++++--------- .../interoperability/_iop/operators.py | 10 ++--- .../interoperability/_iop/reductions.py | 14 +++---- decent_array/interoperability/_iop/rng.py | 18 ++++----- decent_array/interoperability/_iop/utils.py | 10 ++--- .../interoperability/_jax/jax_backend.py | 28 +++++++++----- .../interoperability/_numpy/numpy_backend.py | 32 +++++++++------- .../_pytorch/pytorch_backend.py | 20 +++++----- .../_tensorflow/tensorflow_backend.py | 37 +++++++++--------- docs/source/user.rst | 29 ++++++++++++++ pyproject.toml | 1 + 18 files changed, 224 insertions(+), 138 deletions(-) create mode 100644 decent_array/_errors.py diff --git a/decent_array/_array.py b/decent_array/_array.py index 861c16f..5580c4a 100644 --- a/decent_array/_array.py +++ b/decent_array/_array.py @@ -451,5 +451,6 @@ def item(self) -> Any: # noqa: ANN401 Raises: TypeError: if ``x`` is not 0-dimensional. - """ # noqa: DOC502 + + """ return self._backend.to_scalar(self) diff --git a/decent_array/_errors.py b/decent_array/_errors.py new file mode 100644 index 0000000..cb6abd7 --- /dev/null +++ b/decent_array/_errors.py @@ -0,0 +1,37 @@ +"""Common errors used across backends.""" + +from typing import Any + +from mypy_extensions import mypyc_attr + + +@mypyc_attr(native_class=False) +class NotScalarError(TypeError): + def __init__(self, ndim: int): + super().__init__(f"Only 0-dim arrays can be converted to Python scalars, got {ndim}-dim array.") + + +class NDimError(ValueError): + def __init__(self, required_ndim: int, actual_ndim: int): + super().__init__(f"A {required_ndim}-dim array is required, got {actual_ndim}-dim array.") + + +class MatrixTransposeError(ValueError): + def __init__(self, ndim: int): + super().__init__(f"An aray with at least 2 dimensions is required, got {ndim}-dim array.") + + +class UnsupportedDTypeCreationError(ValueError): + def __init__(self, dtype: Any, backend_name: str, device_name: str): # noqa: ANN401 + super().__init__(f"Unsupported dtype '{dtype}' for {backend_name} on {device_name}.") + + +class UnsupportedDeviceError(ValueError): + def __init__(self, backend_name: str, device_name: str): + super().__init__(f"{backend_name} does not support device '{device_name}'.") + + +stack_empty_error = ValueError("Cannot stack an empty sequence of arrays.") + + +no_backend_error = RuntimeError("No backend active: call 'set_backend' with a supported framework to activate one.") diff --git a/decent_array/interoperability/_iop/bit_operators.py b/decent_array/interoperability/_iop/bit_operators.py index 48362bd..0bbbfc5 100644 --- a/decent_array/interoperability/_iop/bit_operators.py +++ b/decent_array/interoperability/_iop/bit_operators.py @@ -15,6 +15,7 @@ from typing import TYPE_CHECKING +from decent_array._errors import no_backend_error from decent_array.interoperability._backend_manager import register_backend_listener if TYPE_CHECKING: @@ -36,40 +37,40 @@ def _update_backend(backend: Backend | None) -> None: def bitwise_and(x1: bool | int | Array, x2: bool | int | Array) -> Array: """Element-wise bitwise/logical AND.""" if _BACKEND_INSTANCE is None: - raise _error + raise no_backend_error return _BACKEND_INSTANCE.bitwise_and(x1, x2) def bitwise_invert(x: Array) -> Array: """Element-wise bitwise/logical NOT.""" if _BACKEND_INSTANCE is None: - raise _error + raise no_backend_error return _BACKEND_INSTANCE.bitwise_invert(x) def bitwise_or(x1: bool | int | Array, x2: bool | int | Array) -> Array: """Element-wise bitwise/logical OR.""" if _BACKEND_INSTANCE is None: - raise _error + raise no_backend_error return _BACKEND_INSTANCE.bitwise_or(x1, x2) def bitwise_xor(x1: bool | int | Array, x2: bool | int | Array) -> Array: """Element-wise bitwise/logical XOR.""" if _BACKEND_INSTANCE is None: - raise _error + raise no_backend_error return _BACKEND_INSTANCE.bitwise_xor(x1, x2) def bitwise_left_shift(x1: int | Array, x2: int | Array) -> Array: """Element-wise bitwise left shift.""" if _BACKEND_INSTANCE is None: - raise _error + raise no_backend_error return _BACKEND_INSTANCE.bitwise_left_shift(x1, x2) def bitwise_right_shift(x1: int | Array, x2: int | Array) -> Array: """Element-wise bitwise right shift.""" if _BACKEND_INSTANCE is None: - raise _error + raise no_backend_error return _BACKEND_INSTANCE.bitwise_right_shift(x1, x2) diff --git a/decent_array/interoperability/_iop/comparison.py b/decent_array/interoperability/_iop/comparison.py index 64d176c..5b28f1c 100644 --- a/decent_array/interoperability/_iop/comparison.py +++ b/decent_array/interoperability/_iop/comparison.py @@ -15,6 +15,7 @@ from typing import TYPE_CHECKING +from decent_array._errors import no_backend_error from decent_array.interoperability._backend_manager import register_backend_listener if TYPE_CHECKING: @@ -22,7 +23,6 @@ from decent_array.interoperability._abstracts import Backend _BACKEND_INSTANCE: Backend | None = None -_error = RuntimeError("No backend active: call 'set_backend' with a supported framework to activate one.") def _update_backend(backend: Backend | None) -> None: @@ -36,40 +36,40 @@ def _update_backend(backend: Backend | None) -> None: def equal(x1: int | float | complex | Array, x2: int | float | complex | Array) -> Array: """Element-wise equality. Returns an :class:`~decent_array.Array` of bools.""" if _BACKEND_INSTANCE is None: - raise _error + raise no_backend_error return _BACKEND_INSTANCE.equal(x1, x2) def not_equal(x1: int | float | complex | Array, x2: int | float | complex | Array) -> Array: """Element-wise inequality. Returns an :class:`~decent_array.Array` of bools.""" if _BACKEND_INSTANCE is None: - raise _error + raise no_backend_error return _BACKEND_INSTANCE.not_equal(x1, x2) def less(x1: int | float | complex | Array, x2: int | float | complex | Array) -> Array: """Element-wise less-than. Returns an :class:`~decent_array.Array` of bools.""" if _BACKEND_INSTANCE is None: - raise _error + raise no_backend_error return _BACKEND_INSTANCE.less(x1, x2) def less_equal(x1: int | float | complex | Array, x2: int | float | complex | Array) -> Array: """Element-wise less-than-or-equal. Returns an :class:`~decent_array.Array` of bools.""" if _BACKEND_INSTANCE is None: - raise _error + raise no_backend_error return _BACKEND_INSTANCE.less_equal(x1, x2) def greater(x1: int | float | complex | Array, x2: int | float | complex | Array) -> Array: """Element-wise greater-than. Returns an :class:`~decent_array.Array` of bools.""" if _BACKEND_INSTANCE is None: - raise _error + raise no_backend_error return _BACKEND_INSTANCE.greater(x1, x2) def greater_equal(x1: int | float | complex | Array, x2: int | float | complex | Array) -> Array: """Element-wise greater-than-or-equal. Returns an :class:`~decent_array.Array` of bools.""" if _BACKEND_INSTANCE is None: - raise _error + raise no_backend_error return _BACKEND_INSTANCE.greater_equal(x1, x2) diff --git a/decent_array/interoperability/_iop/creation.py b/decent_array/interoperability/_iop/creation.py index 1ebcc74..b849204 100644 --- a/decent_array/interoperability/_iop/creation.py +++ b/decent_array/interoperability/_iop/creation.py @@ -15,6 +15,7 @@ from typing import TYPE_CHECKING +from decent_array._errors import no_backend_error from decent_array.interoperability._backend_manager import register_backend_listener if TYPE_CHECKING: @@ -22,7 +23,6 @@ from decent_array.interoperability._abstracts import Backend _BACKEND_INSTANCE: Backend | None = None -_error = RuntimeError("No backend active: call 'set_backend' with a supported framework to activate one.") def _update_backend(backend: Backend | None) -> None: @@ -36,33 +36,33 @@ def _update_backend(backend: Backend | None) -> None: def zeros(shape: int | tuple[int, ...]) -> Array: """Create an array of zeros with the given shape.""" if _BACKEND_INSTANCE is None: - raise _error + raise no_backend_error return _BACKEND_INSTANCE.zeros(shape) def zeros_like(x: Array) -> Array: """Create an array of zeros matching the shape and type of ``x``.""" if _BACKEND_INSTANCE is None: - raise _error + raise no_backend_error return _BACKEND_INSTANCE.zeros_like(x) def ones(shape: int | tuple[int, ...]) -> Array: """Create an array of ones with the given shape.""" if _BACKEND_INSTANCE is None: - raise _error + raise no_backend_error return _BACKEND_INSTANCE.ones(shape) def ones_like(x: Array) -> Array: """Create an array of ones matching the shape and type of ``x``.""" if _BACKEND_INSTANCE is None: - raise _error + raise no_backend_error return _BACKEND_INSTANCE.ones_like(x) def eye(n: int) -> Array: """Create an ``n x n`` identity matrix.""" if _BACKEND_INSTANCE is None: - raise _error + raise no_backend_error return _BACKEND_INSTANCE.eye(n) diff --git a/decent_array/interoperability/_iop/linalg.py b/decent_array/interoperability/_iop/linalg.py index c10e0f0..44b1792 100644 --- a/decent_array/interoperability/_iop/linalg.py +++ b/decent_array/interoperability/_iop/linalg.py @@ -15,6 +15,7 @@ from typing import TYPE_CHECKING +from decent_array._errors import no_backend_error from decent_array.interoperability._backend_manager import register_backend_listener if TYPE_CHECKING: @@ -22,7 +23,6 @@ from decent_array.interoperability._abstracts import Backend _BACKEND_INSTANCE: Backend | None = None -_error = RuntimeError("No backend active: call 'set_backend' with a supported framework to activate one.") def _update_backend(backend: Backend | None) -> None: @@ -36,7 +36,7 @@ def _update_backend(backend: Backend | None) -> None: def vecdot(x1: Array, x2: Array) -> Array: """Vector dot product of two arrays.""" if _BACKEND_INSTANCE is None: - raise _error + raise no_backend_error return _BACKEND_INSTANCE.vecdot(x1, x2) @@ -47,14 +47,14 @@ def dot(x1: Array, x2: Array) -> Array: Alias for :func:`vecdot`. """ if _BACKEND_INSTANCE is None: - raise _error + raise no_backend_error return _BACKEND_INSTANCE.vecdot(x1, x2) def matmul(x1: Array, x2: Array) -> Array: """Matrix multiplication of two arrays.""" if _BACKEND_INSTANCE is None: - raise _error + raise no_backend_error return _BACKEND_INSTANCE.matmul(x1, x2) @@ -66,7 +66,7 @@ def vector_norm( ) -> Array: """Vector norm of ``x``.""" if _BACKEND_INSTANCE is None: - raise _error + raise no_backend_error return _BACKEND_INSTANCE.vector_norm(x, axis, keepdims, ord) @@ -82,5 +82,5 @@ def norm( Alias for :func:`vector_norm`. """ if _BACKEND_INSTANCE is None: - raise _error + raise no_backend_error return _BACKEND_INSTANCE.vector_norm(x, axis, keepdims, ord) diff --git a/decent_array/interoperability/_iop/manipulations.py b/decent_array/interoperability/_iop/manipulations.py index 629a2f8..e04bf28 100644 --- a/decent_array/interoperability/_iop/manipulations.py +++ b/decent_array/interoperability/_iop/manipulations.py @@ -16,6 +16,7 @@ from collections.abc import Sequence from typing import TYPE_CHECKING, Any +from decent_array._errors import no_backend_error from decent_array.interoperability._backend_manager import register_backend_listener if TYPE_CHECKING: @@ -27,7 +28,6 @@ from decent_array.types._dtypes import dtype _BACKEND_INSTANCE: Backend | None = None -_error = RuntimeError("No backend active: call 'set_backend' with a supported framework to activate one.") def _update_backend(backend: Backend | None) -> None: @@ -41,124 +41,124 @@ def _update_backend(backend: Backend | None) -> None: def copy(x: Array) -> Array: """Return a copy of ``x``.""" if _BACKEND_INSTANCE is None: - raise _error + raise no_backend_error return _BACKEND_INSTANCE.copy(x) def to_numpy(x: ArrayTypes | Array) -> NDArray[Any]: """Convert ``x`` to a NumPy array on CPU.""" if _BACKEND_INSTANCE is None: - raise _error + raise no_backend_error return _BACKEND_INSTANCE.to_numpy(x) def from_numpy(x: NDArray[Any]) -> Array: """Convert a NumPy array on CPU to an :class:`~decent_array.Array` on the active backend.""" if _BACKEND_INSTANCE is None: - raise _error + raise no_backend_error return _BACKEND_INSTANCE.from_numpy(x) def from_numpy_like(x: NDArray[Any], like: Array) -> Array: """Convert a NumPy array to an :class:`~decent_array.Array` matching ``like``'s dtype and device.""" if _BACKEND_INSTANCE is None: - raise _error + raise no_backend_error return _BACKEND_INSTANCE.from_numpy_like(x, like) def asarray(x: float | bool) -> Array: """Convert a Python scalar to an :class:`~decent_array.Array` on the active backend.""" if _BACKEND_INSTANCE is None: - raise _error + raise no_backend_error return _BACKEND_INSTANCE.asarray(x) def stack(arrays: Sequence[Array], axis: int = 0) -> Array: """Stack a sequence of arrays along a new dimension.""" if _BACKEND_INSTANCE is None: - raise _error + raise no_backend_error return _BACKEND_INSTANCE.stack(arrays, axis) def reshape(x: Array, shape: tuple[int, ...]) -> Array: """Reshape ``x`` to ``shape``.""" if _BACKEND_INSTANCE is None: - raise _error + raise no_backend_error return _BACKEND_INSTANCE.reshape(x, shape) def transpose(x: Array, axis: tuple[int, ...] | None = None) -> Array: """Transpose ``x``; ``None`` reverses dimensions.""" if _BACKEND_INSTANCE is None: - raise _error + raise no_backend_error return _BACKEND_INSTANCE.transpose(x, axis) def matrix_transpose(x: Array) -> Array: """Transpose the innermost two dimensions of ``x``.""" if _BACKEND_INSTANCE is None: - raise _error + raise no_backend_error return _BACKEND_INSTANCE.matrix_transpose(x) def shape(x: Array) -> tuple[int, ...]: """Return the shape of ``x``.""" if _BACKEND_INSTANCE is None: - raise _error + raise no_backend_error return _BACKEND_INSTANCE.shape(x) def size(x: Array) -> int: """Return the total number of elements in ``x``.""" if _BACKEND_INSTANCE is None: - raise _error + raise no_backend_error return _BACKEND_INSTANCE.size(x) def ndim(x: Array) -> int: """Return the number of dimensions of ``x``.""" if _BACKEND_INSTANCE is None: - raise _error + raise no_backend_error return _BACKEND_INSTANCE.ndim(x) def squeeze(x: Array, axis: int | tuple[int, ...] | None = None) -> Array: """Remove single-dimensional entries from ``x``.""" if _BACKEND_INSTANCE is None: - raise _error + raise no_backend_error return _BACKEND_INSTANCE.squeeze(x, axis) def expand_dims(x: Array, axis: int) -> Array: """Insert a singleton dimension at ``axis``.""" if _BACKEND_INSTANCE is None: - raise _error + raise no_backend_error return _BACKEND_INSTANCE.unsqueeze(x, axis) def unsqueeze(x: Array, axis: int) -> Array: """Insert a singleton dimension at ``axis``.""" if _BACKEND_INSTANCE is None: - raise _error + raise no_backend_error return _BACKEND_INSTANCE.unsqueeze(x, axis) def diag(x: Array) -> Array: """Build a diagonal matrix from a 1-D vector.""" if _BACKEND_INSTANCE is None: - raise _error + raise no_backend_error return _BACKEND_INSTANCE.diag(x) def diagonal(x: Array, offset: int = 0) -> Array: """Extract the diagonal entries from a 2-D matrix at the given ``offset``.""" if _BACKEND_INSTANCE is None: - raise _error + raise no_backend_error return _BACKEND_INSTANCE.diagonal(x, offset) def astype(x: Array, dtype: dtype) -> Array: """Cast ``x`` to a different dtype.""" if _BACKEND_INSTANCE is None: - raise _error + raise no_backend_error return _BACKEND_INSTANCE.astype(x, dtype) diff --git a/decent_array/interoperability/_iop/math.py b/decent_array/interoperability/_iop/math.py index 1b6bcc7..5ff6f6c 100644 --- a/decent_array/interoperability/_iop/math.py +++ b/decent_array/interoperability/_iop/math.py @@ -15,6 +15,7 @@ from typing import TYPE_CHECKING +from decent_array._errors import no_backend_error from decent_array.interoperability._backend_manager import register_backend_listener if TYPE_CHECKING: @@ -22,7 +23,6 @@ from decent_array.interoperability._abstracts import Backend _BACKEND_INSTANCE: Backend | None = None -_error = RuntimeError("No backend active: call 'set_backend' with a supported framework to activate one.") def _update_backend(backend: Backend | None) -> None: @@ -36,98 +36,98 @@ def _update_backend(backend: Backend | None) -> None: def add(x1: Array | float, x2: Array | float) -> Array: """Element-wise addition.""" if _BACKEND_INSTANCE is None: - raise _error + raise no_backend_error return _BACKEND_INSTANCE.add(x1, x2) def iadd[T: Array](x1: T, x2: Array | float) -> T: """In-place element-wise addition.""" if _BACKEND_INSTANCE is None: - raise _error + raise no_backend_error return _BACKEND_INSTANCE.iadd(x1, x2) def subtract(x1: int | float | complex | Array, x2: int | float | complex | Array) -> Array: """Element-wise subtraction.""" if _BACKEND_INSTANCE is None: - raise _error + raise no_backend_error return _BACKEND_INSTANCE.subtract(x1, x2) def isubtract[T: Array](x1: T, x2: int | float | complex | Array) -> T: """In-place element-wise subtraction.""" if _BACKEND_INSTANCE is None: - raise _error + raise no_backend_error return _BACKEND_INSTANCE.isubtract(x1, x2) def multiply(x1: Array | float, x2: int | float | complex | Array) -> Array: """Element-wise multiplication.""" if _BACKEND_INSTANCE is None: - raise _error + raise no_backend_error return _BACKEND_INSTANCE.multiply(x1, x2) def imultiply[T: Array](x1: T, x2: int | float | complex | Array) -> T: """In-place element-wise multiplication.""" if _BACKEND_INSTANCE is None: - raise _error + raise no_backend_error return _BACKEND_INSTANCE.imultiply(x1, x2) def divide(x1: Array | float, x2: int | float | complex | Array) -> Array: """Element-wise division.""" if _BACKEND_INSTANCE is None: - raise _error + raise no_backend_error return _BACKEND_INSTANCE.divide(x1, x2) def idivide[T: Array](x1: T, x2: int | float | complex | Array) -> T: """In-place element-wise division.""" if _BACKEND_INSTANCE is None: - raise _error + raise no_backend_error return _BACKEND_INSTANCE.idivide(x1, x2) def floor_divide(x1: int | float | Array, x2: int | float | Array) -> Array: """Element-wise floor division.""" if _BACKEND_INSTANCE is None: - raise _error + raise no_backend_error return _BACKEND_INSTANCE.floor_divide(x1, x2) def remainder(x1: int | float | Array, x2: int | float | Array) -> Array: """Element-wise remainder after floor division.""" if _BACKEND_INSTANCE is None: - raise _error + raise no_backend_error return _BACKEND_INSTANCE.remainder(x1, x2) def pow(x1: int | float | complex | Array, x2: int | float | complex | Array) -> Array: # noqa: A001 """Raise ``x`` to power ``p``.""" if _BACKEND_INSTANCE is None: - raise _error + raise no_backend_error return _BACKEND_INSTANCE.pow(x1, x2) def negative(x: Array) -> Array: """Element-wise negation.""" if _BACKEND_INSTANCE is None: - raise _error + raise no_backend_error return _BACKEND_INSTANCE.negative(x) def positive(x: Array) -> Array: """Return the array itself.""" if _BACKEND_INSTANCE is None: - raise _error + raise no_backend_error return x def absolute(x: Array) -> Array: """Element-wise absolute value.""" if _BACKEND_INSTANCE is None: - raise _error + raise no_backend_error return _BACKEND_INSTANCE.absolute(x) @@ -138,12 +138,12 @@ def abs(x: Array) -> Array: # noqa: A001 Alias for :func:`absolute`. """ if _BACKEND_INSTANCE is None: - raise _error + raise no_backend_error return _BACKEND_INSTANCE.absolute(x) def sqrt(x: Array) -> Array: """Element-wise square root.""" if _BACKEND_INSTANCE is None: - raise _error + raise no_backend_error return _BACKEND_INSTANCE.sqrt(x) diff --git a/decent_array/interoperability/_iop/operators.py b/decent_array/interoperability/_iop/operators.py index b72c6ab..8843edd 100644 --- a/decent_array/interoperability/_iop/operators.py +++ b/decent_array/interoperability/_iop/operators.py @@ -15,6 +15,7 @@ from typing import TYPE_CHECKING +from decent_array._errors import no_backend_error from decent_array.interoperability._backend_manager import register_backend_listener if TYPE_CHECKING: @@ -22,7 +23,6 @@ from decent_array.interoperability._abstracts import Backend _BACKEND_INSTANCE: Backend | None = None -_error = RuntimeError("No backend active: call 'set_backend' with a supported framework to activate one.") def _update_backend(backend: Backend | None) -> None: @@ -36,26 +36,26 @@ def _update_backend(backend: Backend | None) -> None: def sign(x: Array) -> Array: """Element-wise sign.""" if _BACKEND_INSTANCE is None: - raise _error + raise no_backend_error return _BACKEND_INSTANCE.sign(x) def maximum(x1: Array | float, x2: Array | float) -> Array: """Element-wise maximum.""" if _BACKEND_INSTANCE is None: - raise _error + raise no_backend_error return _BACKEND_INSTANCE.maximum(x1, x2) def argmax(x: Array, axis: int | None = None, keepdims: bool = False) -> Array: """Index of maximum value along ``axis``.""" if _BACKEND_INSTANCE is None: - raise _error + raise no_backend_error return _BACKEND_INSTANCE.argmax(x, axis, keepdims) def argmin(x: Array, axis: int | None = None, keepdims: bool = False) -> Array: """Index of minimum value along ``axis``.""" if _BACKEND_INSTANCE is None: - raise _error + raise no_backend_error return _BACKEND_INSTANCE.argmin(x, axis, keepdims) diff --git a/decent_array/interoperability/_iop/reductions.py b/decent_array/interoperability/_iop/reductions.py index 7c81125..20540cc 100644 --- a/decent_array/interoperability/_iop/reductions.py +++ b/decent_array/interoperability/_iop/reductions.py @@ -15,6 +15,7 @@ from typing import TYPE_CHECKING +from decent_array._errors import no_backend_error from decent_array.interoperability._backend_manager import register_backend_listener if TYPE_CHECKING: @@ -22,7 +23,6 @@ from decent_array.interoperability._abstracts import Backend _BACKEND_INSTANCE: Backend | None = None -_error = RuntimeError("No backend active: call 'set_backend' with a supported framework to activate one.") def _update_backend(backend: Backend | None) -> None: @@ -40,7 +40,7 @@ def sum( # noqa: A001 ) -> Array: """Sum elements of ``x`` along ``axis``.""" if _BACKEND_INSTANCE is None: - raise _error + raise no_backend_error return _BACKEND_INSTANCE.sum(x, axis, keepdims) @@ -51,7 +51,7 @@ def mean( ) -> Array: """Mean of ``x`` along ``axis``.""" if _BACKEND_INSTANCE is None: - raise _error + raise no_backend_error return _BACKEND_INSTANCE.mean(x, axis, keepdims) @@ -62,7 +62,7 @@ def min( # noqa: A001 ) -> Array: """Minimum of ``x`` along ``axis``.""" if _BACKEND_INSTANCE is None: - raise _error + raise no_backend_error return _BACKEND_INSTANCE.min(x, axis, keepdims) @@ -73,19 +73,19 @@ def max( # noqa: A001 ) -> Array: """Maximum of ``x`` along ``axis``.""" if _BACKEND_INSTANCE is None: - raise _error + raise no_backend_error return _BACKEND_INSTANCE.max(x, axis, keepdims) def any(x: Array, axis: int | tuple[int, ...] | None = None, keepdims: bool = False) -> bool: # noqa: A001 """Return True if any element of ``x`` is truthy.""" if _BACKEND_INSTANCE is None: - raise _error + raise no_backend_error return _BACKEND_INSTANCE.any(x, axis, keepdims) def all(x: Array, axis: int | tuple[int, ...] | None = None, keepdims: bool = False) -> bool: # noqa: A001 """Return True if all elements of ``x`` are truthy.""" if _BACKEND_INSTANCE is None: - raise _error + raise no_backend_error return _BACKEND_INSTANCE.all(x, axis, keepdims) diff --git a/decent_array/interoperability/_iop/rng.py b/decent_array/interoperability/_iop/rng.py index 3df90ca..7b39029 100644 --- a/decent_array/interoperability/_iop/rng.py +++ b/decent_array/interoperability/_iop/rng.py @@ -20,6 +20,7 @@ import random from typing import TYPE_CHECKING, Any, cast +from decent_array._errors import no_backend_error from decent_array.interoperability._backend_manager import _instantiate, register_backend_listener from decent_array.types import Devices, Frameworks @@ -34,7 +35,6 @@ _NUMPY_STATE_KEY = "__numpy_rng_state__" _PYTHON_RANDOM_KEY = "__python_random_state__" _BACKEND_INSTANCE: Backend | None = None -_error = RuntimeError("No backend active: call 'set_backend' with a supported framework to activate one.") def _update_backend(backend: Backend | None) -> None: @@ -63,7 +63,7 @@ def set_seed(self, seed: int, *, set_global_seed: bool = True) -> None: """ if _BACKEND_INSTANCE is None: - raise _error + raise no_backend_error random.seed(seed) active = _BACKEND_INSTANCE @@ -89,7 +89,7 @@ def get_rng_state(self) -> dict[str, Any]: """ if _BACKEND_INSTANCE is None: - raise _error + raise no_backend_error active = _BACKEND_INSTANCE state = active.get_rng_state() @@ -102,7 +102,7 @@ def get_rng_state(self) -> dict[str, Any]: def set_rng_state(self, state: dict[str, Any]) -> None: """Restore a snapshot produced by :meth:`get_rng_state`.""" if _BACKEND_INSTANCE is None: - raise _error + raise no_backend_error # Copy so we can mutate without surprising the caller. state = dict(state) @@ -188,33 +188,33 @@ def derive_seed() -> int: def normal(mean: float = 0.0, std: float = 1.0, shape: tuple[int, ...] = ()) -> Array: """Draw normally distributed samples on the active backend.""" if _BACKEND_INSTANCE is None: - raise _error + raise no_backend_error return _BACKEND_INSTANCE.normal(mean, std, shape) def uniform(low: float = 0.0, high: float = 1.0, shape: tuple[int, ...] = ()) -> Array: """Draw uniformly distributed samples on the active backend.""" if _BACKEND_INSTANCE is None: - raise _error + raise no_backend_error return _BACKEND_INSTANCE.uniform(low, high, shape) def normal_like(x: Array, mean: float = 0.0, std: float = 1.0) -> Array: """Draw normally distributed samples shaped like ``x`` with same dtype.""" if _BACKEND_INSTANCE is None: - raise _error + raise no_backend_error return _BACKEND_INSTANCE.normal_like(x, mean, std) def uniform_like(x: Array, low: float = 0.0, high: float = 1.0) -> Array: """Draw uniformly distributed samples shaped like ``x`` with same dtype.""" if _BACKEND_INSTANCE is None: - raise _error + raise no_backend_error return _BACKEND_INSTANCE.uniform_like(x, low, high) def choice(x: Array, size: int, replace: bool = True) -> Array: """Sample ``size`` elements from ``x``.""" if _BACKEND_INSTANCE is None: - raise _error + raise no_backend_error return _BACKEND_INSTANCE.choice(x, size, replace) diff --git a/decent_array/interoperability/_iop/utils.py b/decent_array/interoperability/_iop/utils.py index ab69d22..d4014ea 100644 --- a/decent_array/interoperability/_iop/utils.py +++ b/decent_array/interoperability/_iop/utils.py @@ -15,6 +15,7 @@ from typing import TYPE_CHECKING, Any +from decent_array._errors import no_backend_error from decent_array.interoperability._backend_manager import register_backend_listener if TYPE_CHECKING: @@ -23,7 +24,6 @@ from decent_array.types import ArrayKey, Devices _BACKEND_INSTANCE: Backend | None = None -_error = RuntimeError("No backend active: call 'set_backend' with a supported framework to activate one.") def _update_backend(backend: Backend | None) -> None: @@ -37,26 +37,26 @@ def _update_backend(backend: Backend | None) -> None: def device_to_native(device: Devices) -> Any: # noqa: ANN401 """Convert :class:`~decent_array.types.Devices` to the active backend's native device.""" if _BACKEND_INSTANCE is None: - raise _error + raise no_backend_error return _BACKEND_INSTANCE.device_to_native(device) def device_of(x: Array) -> Devices: """Return the :class:`~decent_array.types.Devices` of ``x``.""" if _BACKEND_INSTANCE is None: - raise _error + raise no_backend_error return _BACKEND_INSTANCE.device_of(x) def set_item(x: Array, key: ArrayKey, value: bool | int | float | complex | Array) -> None: """Set ``x[key] = value`` in place.""" if _BACKEND_INSTANCE is None: - raise _error + raise no_backend_error _BACKEND_INSTANCE.set_item(x, key, value) def get_item(x: Array, key: ArrayKey) -> Array: """Return ``x[key]``.""" if _BACKEND_INSTANCE is None: - raise _error + raise no_backend_error return _BACKEND_INSTANCE.get_item(x, key) diff --git a/decent_array/interoperability/_jax/jax_backend.py b/decent_array/interoperability/_jax/jax_backend.py index 3859b6c..bdc40eb 100644 --- a/decent_array/interoperability/_jax/jax_backend.py +++ b/decent_array/interoperability/_jax/jax_backend.py @@ -20,6 +20,13 @@ from numpy.typing import NDArray from decent_array import Array +from decent_array._errors import ( + MatrixTransposeError, + NDimError, + UnsupportedDeviceError, + UnsupportedDTypeCreationError, + stack_empty_error, +) from decent_array._utils import is_scalar, unwrap from decent_array.interoperability._abstracts import Backend from decent_array.interoperability._backend_manager import register_backend @@ -32,6 +39,8 @@ class JaxBackend(Backend): def __init__(self, device: Devices = Devices.CPU) -> None: super().__init__(device, name=Frameworks.JAX.value) + if device == Devices.MPS: + UnsupportedDeviceError(self.name, device.value) self._native_device: jax.Device = self.device_to_native(device) self._key: jax.Array = jax.random.key(time_ns()) @@ -99,7 +108,7 @@ def to_scalar(self, x: Array) -> Any: # noqa: ANN401 def stack(self, arrays: Sequence[Array], axis: int = 0) -> Array: if len(arrays) == 0: - raise ValueError("Cannot stack an empty sequence of arrays.") + raise stack_empty_error return Array(jnp.stack([a.value for a in arrays], axis=axis)) def reshape(self, x: Array, shape: tuple[int, ...]) -> Array: @@ -109,10 +118,9 @@ def transpose(self, x: Array, axis: tuple[int, ...] | None = None) -> Array: return Array(jnp.transpose(x.value, axes=axis)) def matrix_transpose(self, x: Array) -> Array: - v = x.value - if v.ndim < 2: - raise ValueError(f"matrix_transpose requires an array with at least 2 dimensions, got {v.ndim}-D") - return Array(jnp.swapaxes(v, -1, -2)) + if x.ndim < 2: + raise MatrixTransposeError(x.ndim) + return Array(jnp.swapaxes(x.value, -1, -2)) def shape(self, x: Array) -> tuple[int, ...]: return tuple(x.value.shape) @@ -130,18 +138,18 @@ def unsqueeze(self, x: Array, axis: int) -> Array: return Array(jnp.expand_dims(x.value, axis=axis)) def diag(self, x: Array) -> Array: - if x.value.ndim != 1: - raise ValueError(f"diag requires a 1-D array, got {x.value.ndim}-D") + if x.ndim != 1: + raise NDimError(1, x.ndim) return Array(jnp.diag(x.value)) def diagonal(self, x: Array, offset: int = 0) -> Array: - if x.value.ndim != 2: - raise ValueError(f"diagonal requires a 2-D array, got {x.value.ndim}-D") + if x.ndim != 2: + raise NDimError(2, x.ndim) return Array(jnp.diagonal(x.value, offset=offset)) def astype(self, x: Array, dtype: dtype) -> Array: if not dtype.available: - raise ValueError(f"Unsupported dtype '{dtype}' for JAX backend.") + raise UnsupportedDTypeCreationError(dtype, self.name, self.device.value) return Array(jnp.asarray(x.value, dtype=dtype.backend_dtype)) # Linalg diff --git a/decent_array/interoperability/_numpy/numpy_backend.py b/decent_array/interoperability/_numpy/numpy_backend.py index d032ef4..9a6577b 100644 --- a/decent_array/interoperability/_numpy/numpy_backend.py +++ b/decent_array/interoperability/_numpy/numpy_backend.py @@ -15,6 +15,13 @@ from numpy.typing import NDArray from decent_array import Array +from decent_array._errors import ( + MatrixTransposeError, + NDimError, + UnsupportedDeviceError, + UnsupportedDTypeCreationError, + stack_empty_error, +) from decent_array._utils import is_scalar, unwrap from decent_array.interoperability._abstracts import Backend from decent_array.interoperability._backend_manager import register_backend @@ -26,9 +33,9 @@ class NumpyBackend(Backend): """NumPy implementation of :class:`Backend`.""" def __init__(self, device: Devices = Devices.CPU) -> None: - if device != Devices.CPU: - raise ValueError(f"NumPy backend only supports CPU, got '{device.value}'.") super().__init__(device, name=Frameworks.NUMPY.value) + if device != Devices.CPU: + UnsupportedDeviceError(self.name, device.value) self._rng: np.random.Generator = np.random.default_rng() # Array creation @@ -87,14 +94,14 @@ def to_scalar(self, x: Array) -> Any: # noqa: ANN401 Raises: TypeError: if ``x`` is not 0-dimensional. - """ # noqa: DOC501, DOC502 + """ if not is_scalar(x): raise TypeError("Only 0-dim arrays can be converted to Python scalars.") return x.value.item() def stack(self, arrays: Sequence[Array], axis: int = 0) -> Array: if len(arrays) == 0: - raise ValueError("Cannot stack an empty sequence of arrays.") + raise stack_empty_error return Array(np.stack([a.value for a in arrays], axis=axis)) def reshape(self, x: Array, shape: tuple[int, ...]) -> Array: @@ -104,10 +111,9 @@ def transpose(self, x: Array, axis: tuple[int, ...] | None = None) -> Array: return Array(np.transpose(x.value, axes=axis)) def matrix_transpose(self, x: Array) -> Array: - v = x.value - if v.ndim < 2: - raise ValueError(f"matrix_transpose requires an array with at least 2 dimensions, got {v.ndim}-D") - return Array(np.swapaxes(v, -1, -2)) + if x.ndim < 2: + raise MatrixTransposeError(x.ndim) + return Array(np.swapaxes(x.value, -1, -2)) def shape(self, x: Array) -> tuple[int, ...]: return tuple(x.value.shape) @@ -125,18 +131,18 @@ def unsqueeze(self, x: Array, axis: int) -> Array: return Array(np.expand_dims(x.value, axis=axis)) def diag(self, x: Array) -> Array: - if x.value.ndim != 1: - raise ValueError(f"diag requires a 1-D array, got {x.value.ndim}-D") + if x.ndim != 1: + raise NDimError(1, x.ndim) return Array(np.diag(x.value)) def diagonal(self, x: Array, offset: int = 0) -> Array: - if x.value.ndim != 2: - raise ValueError(f"diagonal requires a 2-D array, got {x.value.ndim}-D") + if x.ndim != 2: + raise NDimError(2, x.ndim) return Array(np.diagonal(x.value, offset=offset)) def astype(self, x: Array, dtype: dtype) -> Array: if not dtype.available: - raise ValueError(f"Unsupported dtype '{dtype}' for NumPy backend.") + raise UnsupportedDTypeCreationError(dtype, self.name, self.device.value) return Array(np.asarray(x.value, dtype=dtype.backend_dtype)) # Linalg diff --git a/decent_array/interoperability/_pytorch/pytorch_backend.py b/decent_array/interoperability/_pytorch/pytorch_backend.py index 7d39174..542a534 100644 --- a/decent_array/interoperability/_pytorch/pytorch_backend.py +++ b/decent_array/interoperability/_pytorch/pytorch_backend.py @@ -15,6 +15,7 @@ from numpy.typing import NDArray from decent_array import Array +from decent_array._errors import MatrixTransposeError, NDimError, UnsupportedDTypeCreationError, stack_empty_error from decent_array._utils import is_scalar, unwrap from decent_array.interoperability._abstracts import Backend from decent_array.interoperability._backend_manager import register_backend @@ -104,7 +105,7 @@ def to_scalar(self, x: Array) -> Any: # noqa: ANN401 def stack(self, arrays: Sequence[Array], axis: int = 0) -> Array: if len(arrays) == 0: - raise ValueError("Cannot stack an empty sequence of arrays.") + raise stack_empty_error return Array(torch.stack([a.value for a in arrays], dim=axis)) def reshape(self, x: Array, shape: tuple[int, ...]) -> Array: @@ -116,10 +117,9 @@ def transpose(self, x: Array, axis: tuple[int, ...] | None = None) -> Array: return Array(torch.permute(v, dims=dims)) def matrix_transpose(self, x: Array) -> Array: - v = x.value - if v.ndim < 2: - raise ValueError(f"matrix_transpose requires an array with at least 2 dimensions, got {v.ndim}-D") - return Array(v.mT) + if x.ndim < 2: + raise MatrixTransposeError(x.ndim) + return Array(x.value.mT) def shape(self, x: Array) -> tuple[int, ...]: return tuple(x.value.shape) @@ -140,18 +140,18 @@ def unsqueeze(self, x: Array, axis: int) -> Array: return Array(torch.unsqueeze(x.value, dim=axis)) def diag(self, x: Array) -> Array: - if x.value.ndim != 1: - raise ValueError(f"diag requires a 1-D array, got {x.value.ndim}-D") + if x.ndim != 1: + raise NDimError(1, x.ndim) return Array(torch.diag(x.value)) def diagonal(self, x: Array, offset: int = 0) -> Array: - if x.value.ndim != 2: - raise ValueError(f"diagonal requires a 2-D array, got {x.value.ndim}-D") + if x.ndim != 2: + raise NDimError(2, x.ndim) return Array(torch.diagonal(x.value, offset=offset)) def astype(self, x: Array, dtype: dtype) -> Array: if not dtype.available: - raise ValueError(f"Unsupported dtype '{dtype}' for PyTorch backend.") + raise UnsupportedDTypeCreationError(dtype, self.name, self.device.value) return Array(x.value.to(dtype=dtype.backend_dtype)) # Linalg diff --git a/decent_array/interoperability/_tensorflow/tensorflow_backend.py b/decent_array/interoperability/_tensorflow/tensorflow_backend.py index 438d6f1..0bd3f51 100644 --- a/decent_array/interoperability/_tensorflow/tensorflow_backend.py +++ b/decent_array/interoperability/_tensorflow/tensorflow_backend.py @@ -18,6 +18,13 @@ from numpy.typing import NDArray from decent_array import Array +from decent_array._errors import ( + MatrixTransposeError, + NDimError, + UnsupportedDeviceError, + UnsupportedDTypeCreationError, + stack_empty_error, +) from decent_array._utils import is_scalar, unwrap from decent_array.interoperability._abstracts import Backend from decent_array.interoperability._backend_manager import register_backend @@ -30,6 +37,8 @@ class TensorflowBackend(Backend): def __init__(self, device: Devices = Devices.CPU) -> None: super().__init__(device, name=Frameworks.TENSORFLOW.value) + if device == Devices.MPS: + UnsupportedDeviceError(self.name, device.value) self._native_device: str = self.device_to_native(device) self._generator: tf.random.Generator = tf.random.Generator.from_non_deterministic_state(alg="philox") @@ -109,7 +118,7 @@ def to_scalar(self, x: Array) -> Any: # noqa: ANN401 def stack(self, arrays: Sequence[Array], axis: int = 0) -> Array: if len(arrays) == 0: - raise ValueError("Cannot stack an empty sequence of arrays.") + raise stack_empty_error return Array(tf.stack([a.value for a in arrays], axis=axis)) def reshape(self, x: Array, shape: tuple[int, ...]) -> Array: @@ -119,11 +128,9 @@ def transpose(self, x: Array, axis: tuple[int, ...] | None = None) -> Array: return Array(tf.transpose(x.value, perm=axis)) def matrix_transpose(self, x: Array) -> Array: - v = x.value - rank = v.shape.ndims - if rank is not None and rank < 2: - raise ValueError(f"matrix_transpose requires an array with at least 2 dimensions, got {rank}-D") - return Array(tf.linalg.matrix_transpose(v)) + if x.ndim < 2: + raise MatrixTransposeError(x.ndim) + return Array(tf.linalg.matrix_transpose(x.value)) def shape(self, x: Array) -> tuple[int, ...]: return cast("tuple[int, ...]", tuple(x.value.shape)) @@ -141,22 +148,18 @@ def unsqueeze(self, x: Array, axis: int) -> Array: return Array(tf.expand_dims(x.value, axis=axis)) def diag(self, x: Array) -> Array: - v = x.value - rank = v.shape.ndims - if rank != 1: - raise ValueError(f"diag requires a 1-D tensor, got rank {rank}") - return Array(tf.linalg.diag(v)) + if x.ndim != 1: + raise NDimError(1, x.ndim) + return Array(tf.linalg.diag(x.value)) def diagonal(self, x: Array, offset: int = 0) -> Array: - v = x.value - rank = v.shape.ndims - if rank != 2: - raise ValueError(f"diagonal requires a 2-D tensor, got rank {rank}") - return Array(tf.linalg.diag_part(v, k=offset)) + if x.ndim != 2: + raise NDimError(2, x.ndim) + return Array(tf.linalg.diag_part(x.value, k=offset)) def astype(self, x: Array, dtype: dtype) -> Array: if not dtype.available: - raise ValueError(f"Unsupported dtype '{dtype}' for TensorFlow backend.") + raise UnsupportedDTypeCreationError(dtype, self.name, self.device.value) return Array(tf.cast(x.value, dtype=dtype.backend_dtype)) # Linalg diff --git a/docs/source/user.rst b/docs/source/user.rst index f31d0b9..4666a57 100644 --- a/docs/source/user.rst +++ b/docs/source/user.rst @@ -10,6 +10,35 @@ Requires `Python 3.13+ `_ pip install decent-array +Supported devices +----------------- + +.. list-table:: device support across frameworks + :header-rows: 1 + :widths: 22 10 10 10 10 + + * - device + - NumPy + - JAX + - PyTorch + - TensorFlow + * - CPU + - ✓ + - ✓ + - ✓ + - ✓ + * - GPU + - + - ✓ + - ✓ + - ✓ + * - MPS + - + - + - ✓ + - + + Constants --------- diff --git a/pyproject.toml b/pyproject.toml index 609a474..e4ffa90 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -13,6 +13,7 @@ classifiers = [ license = "AGPL-3.0-only" dependencies = [ "numpy>=2,<3", + "mypy_extensions>=1.1.0,<2", ] [project.urls] From 9a84679177090c0941a2838707a54bae0744144a Mon Sep 17 00:00:00 2001 From: nicola-bastianello Date: Mon, 6 Jul 2026 20:36:28 +0200 Subject: [PATCH 21/23] enh: first steps for cupy support --- .../interoperability/_cupy/__init__.py | 5 + .../interoperability/_cupy/cupy_backend.py | 481 ++++++++++++++++++ decent_array/types/_types.py | 1 + docs/source/user.rst | 35 +- pyproject.toml | 18 +- tests/conftest.py | 12 + 6 files changed, 549 insertions(+), 3 deletions(-) create mode 100644 decent_array/interoperability/_cupy/__init__.py create mode 100644 decent_array/interoperability/_cupy/cupy_backend.py diff --git a/decent_array/interoperability/_cupy/__init__.py b/decent_array/interoperability/_cupy/__init__.py new file mode 100644 index 0000000..0fbab1c --- /dev/null +++ b/decent_array/interoperability/_cupy/__init__.py @@ -0,0 +1,5 @@ +"""CuPy backend package; importing it triggers backend registration.""" + +from .cupy_backend import CupyBackend + +__all__ = ["CupyBackend"] diff --git a/decent_array/interoperability/_cupy/cupy_backend.py b/decent_array/interoperability/_cupy/cupy_backend.py new file mode 100644 index 0000000..588ad52 --- /dev/null +++ b/decent_array/interoperability/_cupy/cupy_backend.py @@ -0,0 +1,481 @@ +""" +CuPy backend. + +Importing this module registers the backend via :func:`register_backend`, so the +package can be auto-loaded on the first ``set_backend("cupy")`` call. +""" + +from __future__ import annotations + +from collections.abc import Sequence +from copy import deepcopy +from typing import Any, cast + +import numpy as np +import cupy as cp +from numpy.typing import NDArray + +from decent_array import Array +from decent_array._errors import ( + MatrixTransposeError, + NDimError, + UnsupportedDeviceError, + UnsupportedDTypeCreationError, + stack_empty_error, +) +from decent_array._utils import is_scalar, unwrap +from decent_array.interoperability._abstracts import Backend +from decent_array.interoperability._backend_manager import register_backend +from decent_array.types import ArrayKey, ArrayTypes, Devices, Frameworks +from decent_array.types._dtypes import dtype + + +class CupyBackend(Backend): + """CuPy implementation of :class:`Backend`.""" + + def __init__(self, device: Devices = Devices.GPU) -> None: + super().__init__(device, name=Frameworks.CUPY.value) + if device != Devices.GPU: + UnsupportedDeviceError(self.name, device.value) + self._rng: cp.random.Generator = cp.random.default_rng() + + # Array creation + + def zeros(self, shape: int | tuple[int, ...]) -> Array: + return Array(cp.zeros(shape)) + + def zeros_like(self, x: Array) -> Array: + return Array(cp.zeros_like(x.value)) + + def ones(self, shape: int | tuple[int, ...]) -> Array: + return Array(cp.ones(shape)) + + def ones_like(self, x: Array) -> Array: + return Array(cp.ones_like(x.value)) + + def eye(self, n: int) -> Array: + return Array(cp.eye(n)) + + def device_to_native(self, device: Devices) -> Any: # noqa: ANN401 + # NumPy has no explicit device management; surface the request unchanged. + return device + + def device_of(self, x: Array) -> Devices: # noqa: ARG002 + return Devices.CPU + + # Array manipulation + + def copy(self, x: Array) -> Array: + v = x.value + if isinstance(v, cp.ndarray | cp.generic): + return Array(cp.copy(v)) + return Array(deepcopy(v)) + + def to_numpy(self, x: ArrayTypes | Array) -> NDArray[Any]: + """Return the value of an :class:`Array` as a NumPy array.""" + v = x.value if type(x) is Array else x + if isinstance(v, cp.ndarray): + return cp.asnumpy(v) + return np.asarray(v) + + def from_numpy(self, x: NDArray[Any]) -> Array: + return Array(cp.asarray(x)) + + def from_numpy_like(self, x: NDArray[Any], like: Array) -> Array: + return Array(cp.asarray(x, dtype=like.value.dtype)) + + def asarray(self, x: bool | int | float | complex) -> Array: + return Array(cp.array(x)) + + def to_scalar(self, x: Array) -> Any: # noqa: ANN401 + """ + Convert a 0-dim array to a scalar. + + Raises: + TypeError: if ``x`` is not 0-dimensional. + + """ + if not is_scalar(x): + raise TypeError("Only 0-dim arrays can be converted to Python scalars.") + return x.value.item() + + def stack(self, arrays: Sequence[Array], axis: int = 0) -> Array: + if len(arrays) == 0: + raise stack_empty_error + return Array(np.stack([a.value for a in arrays], axis=axis)) + + def reshape(self, x: Array, shape: tuple[int, ...]) -> Array: + return Array(np.reshape(x.value, shape)) + + def transpose(self, x: Array, axis: tuple[int, ...] | None = None) -> Array: + return Array(np.transpose(x.value, axes=axis)) + + def matrix_transpose(self, x: Array) -> Array: + if x.ndim < 2: + raise MatrixTransposeError(x.ndim) + return Array(np.swapaxes(x.value, -1, -2)) + + def shape(self, x: Array) -> tuple[int, ...]: + return tuple(x.value.shape) + + def size(self, x: Array) -> int: + return int(x.value.size) + + def ndim(self, x: Array) -> int: + return int(x.value.ndim) + + def squeeze(self, x: Array, axis: int | tuple[int, ...] | None = None) -> Array: + return Array(np.squeeze(x.value, axis=axis)) + + def unsqueeze(self, x: Array, axis: int) -> Array: + return Array(np.expand_dims(x.value, axis=axis)) + + def diag(self, x: Array) -> Array: + if x.ndim != 1: + raise NDimError(1, x.ndim) + return Array(np.diag(x.value)) + + def diagonal(self, x: Array, offset: int = 0) -> Array: + if x.ndim != 2: + raise NDimError(2, x.ndim) + return Array(np.diagonal(x.value, offset=offset)) + + def astype(self, x: Array, dtype: dtype) -> Array: + if not dtype.available: + raise UnsupportedDTypeCreationError(dtype, self.name, self.device.value) + return Array(np.asarray(x.value, dtype=dtype.backend_dtype)) + + # Linalg + + def vecdot(self, x1: Array, x2: Array) -> Array: + return Array(np.dot(x1.value, x2.value)) + + def matmul(self, x1: Array, x2: Array) -> Array: + return Array(x1.value @ x2.value) + + def imatmul[T: Array](self, x1: T, x2: Array) -> T: + x1.value @= x2.value + return x1 + + def vector_norm( + self, + x: Array, + axis: int | tuple[int, ...] | None = None, + keepdims: bool = False, + ord: int | float = 2, # noqa: A002 + ) -> Array: + # axis in np.linalg.vector_norm seems to allow for int | tuple[int, ...] | None at runtime, + # but is still typed as int | tuple[int, int] | None, hence the ignore + return Array(np.linalg.vector_norm(x.value, ord=ord, axis=axis, keepdims=keepdims)) # type: ignore[arg-type] + + # Math reductions + + def sum(self, x: Array, axis: int | tuple[int, ...] | None = None, keepdims: bool = False) -> Array: + v = cast("np.ndarray[Any, Any]", x.value) + if keepdims: + return Array(np.sum(v, axis=axis, keepdims=True)) + return Array(np.sum(v, axis=axis)) + + def mean(self, x: Array, axis: int | tuple[int, ...] | None = None, keepdims: bool = False) -> Array: + v = cast("np.ndarray[Any, Any]", x.value) + if keepdims: + return Array(np.mean(v, axis=axis, keepdims=True)) + return Array(np.mean(v, axis=axis)) + + def min(self, x: Array, axis: int | tuple[int, ...] | None = None, keepdims: bool = False) -> Array: + v = cast("np.ndarray[Any, Any]", x.value) + if keepdims: + return Array(np.min(v, axis=axis, keepdims=True)) + return Array(np.min(v, axis=axis)) + + def max(self, x: Array, axis: int | tuple[int, ...] | None = None, keepdims: bool = False) -> Array: + v = cast("np.ndarray[Any, Any]", x.value) + if keepdims: + return Array(np.max(v, axis=axis, keepdims=True)) + return Array(np.max(v, axis=axis)) + + def any(self, x: Array, axis: int | tuple[int, ...] | None = None, keepdims: bool = False) -> bool: + v = cast("np.ndarray[Any, Any]", x.value) + if keepdims: + return bool(np.any(v, axis=axis, keepdims=True)) + return bool(np.any(v, axis=axis)) + + def all(self, x: Array, axis: int | tuple[int, ...] | None = None, keepdims: bool = False) -> bool: + v = cast("np.ndarray[Any, Any]", x.value) + if keepdims: + return bool(np.all(v, axis=axis, keepdims=True)) + return bool(np.all(v, axis=axis)) + + # Math elementwise — operands may be Array or scalar (operator dunders pass either). + # ``Array | float`` covers both: PEP 484's numeric tower implicitly admits ``int``. + + def add(self, x1: int | float | complex | Array, x2: int | float | complex | Array) -> Array: + return Array(np.add(unwrap(x1), unwrap(x2))) + + def iadd[T: Array](self, x1: T, x2: int | float | complex | Array) -> T: + x1.value += unwrap(x2) + return x1 + + def subtract(self, x1: int | float | complex | Array, x2: int | float | complex | Array) -> Array: + return Array(np.subtract(unwrap(x1), unwrap(x2))) + + def isubtract[T: Array](self, x1: T, x2: int | float | complex | Array) -> T: + x1.value -= unwrap(x2) + return x1 + + def multiply(self, x1: int | float | complex | Array, x2: int | float | complex | Array) -> Array: + return Array(np.multiply(unwrap(x1), unwrap(x2))) + + def imultiply[T: Array](self, x1: T, x2: int | float | complex | Array) -> T: + x1.value *= unwrap(x2) + return x1 + + def divide(self, x1: int | float | complex | Array, x2: int | float | complex | Array) -> Array: + return Array(np.divide(unwrap(x1), unwrap(x2))) + + def idivide[T: Array](self, x1: T, x2: int | float | complex | Array) -> T: + x1.value /= unwrap(x2) + return x1 + + def floor_divide(self, x1: int | float | Array, x2: int | float | Array) -> Array: + return Array(np.floor_divide(unwrap(x1), unwrap(x2))) + + def ifloordiv[T: Array](self, x1: T, x2: int | float | Array) -> T: + x1.value //= unwrap(x2) + return x1 + + def remainder(self, x1: int | float | Array, x2: int | float | Array) -> Array: + return Array(np.remainder(unwrap(x1), unwrap(x2))) + + def imod[T: Array](self, x1: T, x2: int | float | Array) -> T: + x1.value %= unwrap(x2) + return x1 + + def pow(self, x1: int | float | complex | Array, x2: int | float | complex | Array) -> Array: + return Array(np.power(unwrap(x1), unwrap(x2))) + + def ipow[T: Array](self, x1: T, x2: int | float | complex | Array) -> T: + x1.value **= unwrap(x2) + return x1 + + def negative(self, x: Array) -> Array: + return Array(np.negative(x.value)) + + def absolute(self, x: Array) -> Array: + return Array(np.abs(x.value)) + + def sqrt(self, x: Array) -> Array: + return Array(np.sqrt(x.value)) + + # Comparisons + + def equal(self, x1: int | float | complex | Array, x2: int | float | complex | Array) -> Array: + return Array(np.equal(unwrap(x1), unwrap(x2))) + + def not_equal(self, x1: int | float | complex | Array, x2: int | float | complex | Array) -> Array: + return Array(np.not_equal(unwrap(x1), unwrap(x2))) + + def less(self, x1: int | float | complex | Array, x2: int | float | complex | Array) -> Array: + return Array(np.less(unwrap(x1), unwrap(x2))) + + def less_equal(self, x1: int | float | complex | Array, x2: int | float | complex | Array) -> Array: + return Array(np.less_equal(unwrap(x1), unwrap(x2))) + + def greater(self, x1: int | float | complex | Array, x2: int | float | complex | Array) -> Array: + return Array(np.greater(unwrap(x1), unwrap(x2))) + + def greater_equal(self, x1: int | float | complex | Array, x2: int | float | complex | Array) -> Array: + return Array(np.greater_equal(unwrap(x1), unwrap(x2))) + + # Bitwise + + def bitwise_and(self, x1: bool | int | Array, x2: bool | int | Array) -> Array: + return Array(np.bitwise_and(unwrap(x1), unwrap(x2))) + + def iand[T: Array](self, x1: T, x2: bool | int | Array) -> T: + x1.value &= unwrap(x2) + return x1 + + def bitwise_invert(self, x: Array) -> Array: + return Array(np.bitwise_not(x.value)) + + def bitwise_or(self, x1: bool | int | Array, x2: bool | int | Array) -> Array: + return Array(np.bitwise_or(unwrap(x1), unwrap(x2))) + + def ior[T: Array](self, x1: T, x2: bool | int | Array) -> T: + x1.value |= unwrap(x2) + return x1 + + def bitwise_xor(self, x1: bool | int | Array, x2: bool | int | Array) -> Array: + return Array(np.bitwise_xor(unwrap(x1), unwrap(x2))) + + def ixor[T: Array](self, x1: T, x2: bool | int | Array) -> T: + x1.value ^= unwrap(x2) + return x1 + + def bitwise_left_shift(self, x1: int | Array, x2: int | Array) -> Array: + return Array(np.left_shift(unwrap(x1), unwrap(x2))) + + def ilshift[T: Array](self, x1: T, x2: int | Array) -> T: + x1.value <<= unwrap(x2) + return x1 + + def bitwise_right_shift(self, x1: int | Array, x2: int | Array) -> Array: + return Array(np.right_shift(unwrap(x1), unwrap(x2))) + + def irshift[T: Array](self, x1: T, x2: int | Array) -> T: + x1.value >>= unwrap(x2) + return x1 + + # Operators + + def sign(self, x: Array) -> Array: + return Array(np.sign(x.value)) + + def maximum(self, x1: int | float | complex | Array, x2: int | float | complex | Array) -> Array: + return Array(np.maximum(unwrap(x1), unwrap(x2))) + + def argmax(self, x: Array, axis: int | None = None, keepdims: bool = False) -> Array: + v = cast("np.ndarray[Any, Any]", x.value) + if keepdims: + return Array(np.argmax(v, axis=axis, keepdims=True)) + return Array(np.argmax(v, axis=axis)) + + def argmin(self, x: Array, axis: int | None = None, keepdims: bool = False) -> Array: + v = cast("np.ndarray[Any, Any]", x.value) + if keepdims: + return Array(np.argmin(v, axis=axis, keepdims=True)) + return Array(np.argmin(v, axis=axis)) + + def set_item(self, x: Array, key: ArrayKey, value: bool | int | float | complex | Array) -> None: + x.value[key] = unwrap(value) + + def get_item(self, x: Array, key: ArrayKey) -> Array: + return Array(x.value[key]) + + # RNG + + def set_seed(self, seed: int) -> None: + # Seed both the legacy global state and our owned Generator. The legacy state is + # important because some downstream libraries (sklearn, pandas) consult it. + np.random.seed(seed) # noqa: NPY002 + self._rng = np.random.default_rng(seed) + + def get_rng_state(self) -> dict[str, Any]: + # ``np.random.get_state()`` returns a tuple by default; ``legacy=False`` returns + # the equivalent dict form, which both matches the surrounding ``dict[str, Any]`` + # value type (so mypyc's strict union narrowing is satisfied) and round-trips + # cleanly through ``np.random.set_state``. + return { + "numpy_bit_generator_state": deepcopy(self._rng.bit_generator.state), + "numpy_legacy_state": np.random.get_state(legacy=False), # noqa: NPY002 + } + + def set_rng_state(self, state: dict[str, Any]) -> None: + if "numpy_bit_generator_state" in state: + self._rng = np.random.default_rng() + self._rng.bit_generator.state = state["numpy_bit_generator_state"] + if "numpy_legacy_state" in state: + np.random.set_state(state["numpy_legacy_state"]) # noqa: NPY002 + + def normal(self, mean: float = 0.0, std: float = 1.0, shape: tuple[int, ...] = ()) -> Array: + return Array(self._rng.normal(loc=mean, scale=std, size=shape)) + + def uniform(self, low: float = 0.0, high: float = 1.0, shape: tuple[int, ...] = ()) -> Array: + return Array(self._rng.uniform(low=low, high=high, size=shape)) + + def normal_like(self, x: Array, mean: float = 0.0, std: float = 1.0) -> Array: + return Array(self._rng.normal(loc=mean, scale=std, size=x.value.shape).astype(x.value.dtype)) + + def uniform_like(self, x: Array, low: float = 0.0, high: float = 1.0) -> Array: + return Array(self._rng.uniform(low=low, high=high, size=x.value.shape).astype(x.value.dtype)) + + def choice(self, x: Array, size: int, replace: bool = True) -> Array: + return Array(self._rng.choice(x.value, size=size, replace=replace)) + + # Dtypes + + @property + def bool_(self) -> Any: # noqa: ANN401 + return cp.dtype(cp.bool_) + + @property + def uint8(self) -> Any: # noqa: ANN401 + return cp.dtype(cp.uint8) + + @property + def uint16(self) -> Any: # noqa: ANN401 + return cp.dtype(cp.uint16) + + @property + def uint32(self) -> Any: # noqa: ANN401 + return cp.dtype(cp.uint32) + + @property + def uint64(self) -> Any: # noqa: ANN401 + return cp.dtype(cp) + + @property + def int8(self) -> Any: # noqa: ANN401 + return cp.dtype(cp.int8) + + @property + def int16(self) -> Any: # noqa: ANN401 + return cp.dtype(cp.int16) + + @property + def int32(self) -> Any: # noqa: ANN401 + return cp.dtype(cp.int32) + + @property + def int64(self) -> Any: # noqa: ANN401 + return cp.dtype(cp.int64) + + @property + def float16(self) -> Any: # noqa: ANN401 + return cp.dtype(cp.float16) + + @property + def bfloat16(self) -> Any: # noqa: ANN401 + return cp.dtype(cp.bfloat16) + + @property + def float32(self) -> Any: # noqa: ANN401 + return cp.dtype(cp.float32) + + @property + def float64(self) -> Any: # noqa: ANN401 + return cp.dtype(cp.float64) + + @property + def complex64(self) -> Any: # noqa: ANN401 + return cp.dtype(cp.complex64) + + @property + def complex128(self) -> Any: # noqa: ANN401 + return cp.dtype(cp.complex128) + + # Constants + + @property + def e(self) -> Any: # noqa: ANN401 + """e = 2.71828...""" # noqa: D403 + return cp.e + + @property + def inf(self) -> Any: # noqa: ANN401 + """Infinity.""" + return cp.inf + + @property + def nan(self) -> Any: # noqa: ANN401 + """Not-a-number.""" + return cp.nan + + @property + def pi(self) -> Any: # noqa: ANN401 + """pi = 3.14159...""" # noqa: D403 + return cp.pi + + +register_backend(Frameworks.CUPY, CupyBackend) diff --git a/decent_array/types/_types.py b/decent_array/types/_types.py index 10c6a1b..d2b8ea3 100644 --- a/decent_array/types/_types.py +++ b/decent_array/types/_types.py @@ -44,6 +44,7 @@ class Frameworks(Enum): PYTORCH = "pytorch" TENSORFLOW = "tensorflow" JAX = "jax" + CUPY = "cupy" class Devices(Enum): diff --git a/docs/source/user.rst b/docs/source/user.rst index 4666a57..5f9c298 100644 --- a/docs/source/user.rst +++ b/docs/source/user.rst @@ -15,28 +15,32 @@ Supported devices .. list-table:: device support across frameworks :header-rows: 1 - :widths: 22 10 10 10 10 + :widths: 22 10 10 10 10 10 * - device - NumPy - JAX - PyTorch - TensorFlow + - CuPy * - CPU - ✓ - ✓ - ✓ - ✓ + - * - GPU - - ✓ - ✓ - ✓ + - ✓ * - MPS - - - ✓ - + - Constants @@ -120,19 +124,21 @@ reliable way is to try operations that involve the dtype and observe if the fram .. list-table:: dtype support across frameworks :header-rows: 1 - :widths: 22 10 10 10 12 36 + :widths: 22 10 10 10 10 10 36 * - dtype - NumPy - JAX - PyTorch - TensorFlow + - CuPy - Notes * - ``bool_`` - ✓ - ✓ - ✓ - ✓ + - ✓ - * - **Integers** - @@ -140,29 +146,34 @@ reliable way is to try operations that involve the dtype and observe if the fram - - - + - * - ``int8`` - ✓ - ✓ - ✓ - ✓ + - ✓ - * - ``int16`` - ✓ - ✓ - ✓ - ✓ + - ✓ - * - ``int32`` - ✓ - ✓ - ✓ - ✓ + - ✓ - * - ``int64`` - ✓ - ⚠️ - ✓ - ✓ + - ✓ - JAX requires ``jax_enable_x64=True`` * - **Unsigned integers** - @@ -175,24 +186,28 @@ reliable way is to try operations that involve the dtype and observe if the fram - ✓ - ✓ - ✓ + - ✓ - * - ``uint16`` - ✓ - ⚠️ - ⚠️ - ✓ + - ✓ - JAX requires ``jax_enable_x64=True``; PyTorch support is limited/experimental * - ``uint32`` - ✓ - ⚠️ - ⚠️ - ✓ + - ✓ - JAX requires ``jax_enable_x64=True``; PyTorch support is limited/experimental * - ``uint64`` - ✓ - ⚠️ - ⚠️ - ✓ + - ✓ - JAX requires ``jax_enable_x64=True``; PyTorch support is limited/experimental * - **Floating point** - @@ -211,24 +226,28 @@ reliable way is to try operations that involve the dtype and observe if the fram - ✓ - ✓ - ✓ + - ✓ - PyTorch support is limited/experimental * - ``float32`` - ✓ - ✓ - ✓ - ✓ + - ✓ - * - ``float64`` - ✓ - ⚠️ - ✓ - ✓ + - ✓ - JAX requires ``jax_enable_x64=True``; PyTorch does not support on MPS * - ``float128`` - ⚠️ - ✗ - ✗ - ✗ + - ✗ - NumPy support is platform-dependent * - **Complex** - @@ -241,18 +260,21 @@ reliable way is to try operations that involve the dtype and observe if the fram - ✓ - ✓ - ✓ + - ✓ - * - ``complex128`` - ✓ - ⚠️ - ✓ - ✓ + - ✓ - JAX requires ``jax_enable_x64=True`` * - ``complex256`` - ⚠️ - ✗ - ✗ - ✗ + - ✗ - NumPy support is platform-dependent * - **Quantized** - @@ -265,30 +287,35 @@ reliable way is to try operations that involve the dtype and observe if the fram - ✗ - ✓ - ✓ + - ✗ - * - ``quint8`` - ✗ - ✗ - ✓ - ✓ + - ✗ - * - ``qint16`` - ✗ - ✗ - ✗ - ✓ + - ✗ - * - ``quint16`` - ✗ - ✗ - ✗ - ✓ + - ✗ - * - ``qint32`` - ✗ - ✗ - ✓ - ✓ + - ✗ - * - **Miscellaneous** - @@ -301,22 +328,26 @@ reliable way is to try operations that involve the dtype and observe if the fram - ✗ - ✗ - ✗ + - ✗ - Equivalent to ``np.str_`` * - ``bytes_`` - ✓ - ✗ - ✗ - ✓ + - ✗ - Equivalent to ``np.bytes_`` and ``tf.string`` * - ``object_`` - ✓ - ✗ - ✗ - ✗ + - ✗ - * - ``void`` - ✓ - ✗ - ✗ - ✗ + - ✗ - \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index e4ffa90..a30005a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -38,6 +38,16 @@ jax = [ jax-gpu = [ "jax[cuda12]>=0.4,<1", ] +cupy-cuda12 = [ + "cupy-cuda12x>=14,<15", +] +cupy-cuda13 = [ + "cupy-cuda13x>=14,<15", +] +cupy-rocm = [ + "cupy-rocm-7-0>=14,<15", +] + [dependency-groups] lint = ["ruff==0.15.20"] @@ -77,6 +87,12 @@ include = [ ] mypy-args = ["--ignore-missing-imports"] +[[tool.mypy.overrides]] +module = "cupy" +ignore_missing_imports = true +follow_imports = "skip" +ignore_errors = true + [tool.tox] envlist = ["dev", "mypy", "mypyc", "pytest", "ruff", "sphinx"] package = "editable" @@ -96,7 +112,7 @@ set_env = { PIP_EXTRA_INDEX_URL = "https://download.pytorch.org/whl/cu128", HATC deps = ["-rdocs/requirements.txt"] commands_pre = [ ["python", "-m", "pip", "install", "--group", "dev"], - ["python", "-m", "pip", "install", ".[torch,tf-gpu,jax-gpu]"], + ["python", "-m", "pip", "install", ".[torch,tf-gpu,jax-gpu,cupy-cuda12,cupy-cuda13,cupy-rocm]"], ] [tool.tox.env.sphinx] diff --git a/tests/conftest.py b/tests/conftest.py index 8e83362..6d2b96c 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -31,6 +31,8 @@ def _framework_importable(framework: Frameworks) -> bool: import jax # noqa: F401, PLC0415 elif framework == Frameworks.TENSORFLOW: import tensorflow # noqa: F401, PLC0415 + elif framework == Frameworks.CUPY: + import cupy # noqa: F401, PLC0415 except ImportError: return False return True @@ -79,6 +81,16 @@ def _device_available(framework: Frameworks, device: Devices) -> bool: return len(tf.config.list_physical_devices("GPU")) > 0 except Exception: return False + if framework == Frameworks.CUPY: + if device != Devices.GPU: + return False + import cupy as cp # noqa: PLC0415 + + if device == Devices.GPU: + try: + return bool(cp.cuda.runtime.getDeviceCount() > 0) + except cp.cuda.runtime.CUDARuntimeError: + return False return False From 2ab147c8dcbb56dc3f8ed404001b61bf3a46639f Mon Sep 17 00:00:00 2001 From: nicola-bastianello Date: Thu, 9 Jul 2026 13:55:03 +0200 Subject: [PATCH 22/23] enh: full cupy backend definition --- .../interoperability/_cupy/cupy_backend.py | 169 ++++++++---------- .../interoperability/_jax/jax_backend.py | 5 +- .../interoperability/_numpy/numpy_backend.py | 3 +- .../_pytorch/pytorch_backend.py | 13 +- .../_tensorflow/tensorflow_backend.py | 5 +- docs/source/user.rst | 1 + 6 files changed, 91 insertions(+), 105 deletions(-) diff --git a/decent_array/interoperability/_cupy/cupy_backend.py b/decent_array/interoperability/_cupy/cupy_backend.py index 588ad52..1939924 100644 --- a/decent_array/interoperability/_cupy/cupy_backend.py +++ b/decent_array/interoperability/_cupy/cupy_backend.py @@ -9,9 +9,8 @@ from collections.abc import Sequence from copy import deepcopy -from typing import Any, cast +from typing import Any -import numpy as np import cupy as cp from numpy.typing import NDArray @@ -19,6 +18,7 @@ from decent_array._errors import ( MatrixTransposeError, NDimError, + NotScalarError, UnsupportedDeviceError, UnsupportedDTypeCreationError, stack_empty_error, @@ -37,7 +37,9 @@ def __init__(self, device: Devices = Devices.GPU) -> None: super().__init__(device, name=Frameworks.CUPY.value) if device != Devices.GPU: UnsupportedDeviceError(self.name, device.value) - self._rng: cp.random.Generator = cp.random.default_rng() + self._native_device: cp.cuda.Device = self.device_to_native(device) + self._seed: int = 0 + self._rng: cp.random.Generator = cp.random.default_rng(seed=self._seed) # Array creation @@ -57,11 +59,14 @@ def eye(self, n: int) -> Array: return Array(cp.eye(n)) def device_to_native(self, device: Devices) -> Any: # noqa: ANN401 - # NumPy has no explicit device management; surface the request unchanged. - return device + if device == Devices.GPU: + # currently decent-array only supports specifying "gpu", not the device index; + # so this returns the current device as selected by cupy + return cp.cuda.Device() + raise UnsupportedDeviceError(self.name, device.value) def device_of(self, x: Array) -> Devices: # noqa: ARG002 - return Devices.CPU + return Devices.GPU # Array manipulation @@ -71,12 +76,12 @@ def copy(self, x: Array) -> Array: return Array(cp.copy(v)) return Array(deepcopy(v)) - def to_numpy(self, x: ArrayTypes | Array) -> NDArray[Any]: + def to_numpy(self, x: ArrayTypes | Array) -> cp.ndarray[Any]: """Return the value of an :class:`Array` as a NumPy array.""" v = x.value if type(x) is Array else x if isinstance(v, cp.ndarray): return cp.asnumpy(v) - return np.asarray(v) + return cp.asarray(v) def from_numpy(self, x: NDArray[Any]) -> Array: return Array(cp.asarray(x)) @@ -96,24 +101,24 @@ def to_scalar(self, x: Array) -> Any: # noqa: ANN401 """ if not is_scalar(x): - raise TypeError("Only 0-dim arrays can be converted to Python scalars.") + raise NotScalarError(x.ndim) return x.value.item() def stack(self, arrays: Sequence[Array], axis: int = 0) -> Array: if len(arrays) == 0: raise stack_empty_error - return Array(np.stack([a.value for a in arrays], axis=axis)) + return Array(cp.stack([a.value for a in arrays], axis=axis)) def reshape(self, x: Array, shape: tuple[int, ...]) -> Array: - return Array(np.reshape(x.value, shape)) + return Array(cp.reshape(x.value, shape)) def transpose(self, x: Array, axis: tuple[int, ...] | None = None) -> Array: - return Array(np.transpose(x.value, axes=axis)) + return Array(cp.transpose(x.value, axes=axis)) def matrix_transpose(self, x: Array) -> Array: if x.ndim < 2: raise MatrixTransposeError(x.ndim) - return Array(np.swapaxes(x.value, -1, -2)) + return Array(cp.matrix_transpose(x.value)) def shape(self, x: Array) -> tuple[int, ...]: return tuple(x.value.shape) @@ -125,30 +130,30 @@ def ndim(self, x: Array) -> int: return int(x.value.ndim) def squeeze(self, x: Array, axis: int | tuple[int, ...] | None = None) -> Array: - return Array(np.squeeze(x.value, axis=axis)) + return Array(cp.squeeze(x.value, axis=axis)) def unsqueeze(self, x: Array, axis: int) -> Array: - return Array(np.expand_dims(x.value, axis=axis)) + return Array(cp.expand_dims(x.value, axis=axis)) def diag(self, x: Array) -> Array: if x.ndim != 1: raise NDimError(1, x.ndim) - return Array(np.diag(x.value)) + return Array(cp.diag(x.value)) def diagonal(self, x: Array, offset: int = 0) -> Array: if x.ndim != 2: raise NDimError(2, x.ndim) - return Array(np.diagonal(x.value, offset=offset)) + return Array(cp.diagonal(x.value, offset=offset)) def astype(self, x: Array, dtype: dtype) -> Array: if not dtype.available: raise UnsupportedDTypeCreationError(dtype, self.name, self.device.value) - return Array(np.asarray(x.value, dtype=dtype.backend_dtype)) + return Array(cp.astype(x.value, dtype=dtype.backend_dtype)) # Linalg def vecdot(self, x1: Array, x2: Array) -> Array: - return Array(np.dot(x1.value, x2.value)) + return Array(cp.dot(x1.value, x2.value)) def matmul(self, x1: Array, x2: Array) -> Array: return Array(x1.value @ x2.value) @@ -164,164 +169,146 @@ def vector_norm( keepdims: bool = False, ord: int | float = 2, # noqa: A002 ) -> Array: - # axis in np.linalg.vector_norm seems to allow for int | tuple[int, ...] | None at runtime, - # but is still typed as int | tuple[int, int] | None, hence the ignore - return Array(np.linalg.vector_norm(x.value, ord=ord, axis=axis, keepdims=keepdims)) # type: ignore[arg-type] + if isinstance(axis, tuple) and len(axis) > 2: + raise ValueError(f"'axis' of length {len(axis)} is currently unsupported by {self.name}") + return Array(cp.linalg.norm(x.value, ord=ord, axis=axis, keepdims=keepdims)) # Math reductions def sum(self, x: Array, axis: int | tuple[int, ...] | None = None, keepdims: bool = False) -> Array: - v = cast("np.ndarray[Any, Any]", x.value) - if keepdims: - return Array(np.sum(v, axis=axis, keepdims=True)) - return Array(np.sum(v, axis=axis)) + return Array(cp.sum(x.value, axis=axis, keepdims=keepdims)) def mean(self, x: Array, axis: int | tuple[int, ...] | None = None, keepdims: bool = False) -> Array: - v = cast("np.ndarray[Any, Any]", x.value) - if keepdims: - return Array(np.mean(v, axis=axis, keepdims=True)) - return Array(np.mean(v, axis=axis)) + return Array(cp.mean(x.value, axis=axis, keepdims=keepdims)) def min(self, x: Array, axis: int | tuple[int, ...] | None = None, keepdims: bool = False) -> Array: - v = cast("np.ndarray[Any, Any]", x.value) - if keepdims: - return Array(np.min(v, axis=axis, keepdims=True)) - return Array(np.min(v, axis=axis)) + return Array(cp.min(x.value, axis=axis, keepdims=keepdims)) def max(self, x: Array, axis: int | tuple[int, ...] | None = None, keepdims: bool = False) -> Array: - v = cast("np.ndarray[Any, Any]", x.value) - if keepdims: - return Array(np.max(v, axis=axis, keepdims=True)) - return Array(np.max(v, axis=axis)) + return Array(cp.max(x.value, axis=axis, keepdims=keepdims)) def any(self, x: Array, axis: int | tuple[int, ...] | None = None, keepdims: bool = False) -> bool: - v = cast("np.ndarray[Any, Any]", x.value) - if keepdims: - return bool(np.any(v, axis=axis, keepdims=True)) - return bool(np.any(v, axis=axis)) + return bool(cp.any(x.value, axis=axis, keepdims=keepdims)) def all(self, x: Array, axis: int | tuple[int, ...] | None = None, keepdims: bool = False) -> bool: - v = cast("np.ndarray[Any, Any]", x.value) - if keepdims: - return bool(np.all(v, axis=axis, keepdims=True)) - return bool(np.all(v, axis=axis)) + return bool(cp.all(x.value, axis=axis, keepdims=keepdims)) # Math elementwise — operands may be Array or scalar (operator dunders pass either). # ``Array | float`` covers both: PEP 484's numeric tower implicitly admits ``int``. def add(self, x1: int | float | complex | Array, x2: int | float | complex | Array) -> Array: - return Array(np.add(unwrap(x1), unwrap(x2))) + return Array(cp.add(unwrap(x1), unwrap(x2))) def iadd[T: Array](self, x1: T, x2: int | float | complex | Array) -> T: x1.value += unwrap(x2) return x1 def subtract(self, x1: int | float | complex | Array, x2: int | float | complex | Array) -> Array: - return Array(np.subtract(unwrap(x1), unwrap(x2))) + return Array(cp.subtract(unwrap(x1), unwrap(x2))) def isubtract[T: Array](self, x1: T, x2: int | float | complex | Array) -> T: x1.value -= unwrap(x2) return x1 def multiply(self, x1: int | float | complex | Array, x2: int | float | complex | Array) -> Array: - return Array(np.multiply(unwrap(x1), unwrap(x2))) + return Array(cp.multiply(unwrap(x1), unwrap(x2))) def imultiply[T: Array](self, x1: T, x2: int | float | complex | Array) -> T: x1.value *= unwrap(x2) return x1 def divide(self, x1: int | float | complex | Array, x2: int | float | complex | Array) -> Array: - return Array(np.divide(unwrap(x1), unwrap(x2))) + return Array(cp.divide(unwrap(x1), unwrap(x2))) def idivide[T: Array](self, x1: T, x2: int | float | complex | Array) -> T: x1.value /= unwrap(x2) return x1 def floor_divide(self, x1: int | float | Array, x2: int | float | Array) -> Array: - return Array(np.floor_divide(unwrap(x1), unwrap(x2))) + return Array(cp.floor_divide(unwrap(x1), unwrap(x2))) def ifloordiv[T: Array](self, x1: T, x2: int | float | Array) -> T: x1.value //= unwrap(x2) return x1 def remainder(self, x1: int | float | Array, x2: int | float | Array) -> Array: - return Array(np.remainder(unwrap(x1), unwrap(x2))) + return Array(cp.remainder(unwrap(x1), unwrap(x2))) def imod[T: Array](self, x1: T, x2: int | float | Array) -> T: x1.value %= unwrap(x2) return x1 def pow(self, x1: int | float | complex | Array, x2: int | float | complex | Array) -> Array: - return Array(np.power(unwrap(x1), unwrap(x2))) + return Array(cp.power(unwrap(x1), unwrap(x2))) def ipow[T: Array](self, x1: T, x2: int | float | complex | Array) -> T: x1.value **= unwrap(x2) return x1 def negative(self, x: Array) -> Array: - return Array(np.negative(x.value)) + return Array(cp.negative(x.value)) def absolute(self, x: Array) -> Array: - return Array(np.abs(x.value)) + return Array(cp.absolute(x.value)) def sqrt(self, x: Array) -> Array: - return Array(np.sqrt(x.value)) + return Array(cp.sqrt(x.value)) # Comparisons def equal(self, x1: int | float | complex | Array, x2: int | float | complex | Array) -> Array: - return Array(np.equal(unwrap(x1), unwrap(x2))) + return Array(cp.equal(unwrap(x1), unwrap(x2))) def not_equal(self, x1: int | float | complex | Array, x2: int | float | complex | Array) -> Array: - return Array(np.not_equal(unwrap(x1), unwrap(x2))) + return Array(cp.not_equal(unwrap(x1), unwrap(x2))) def less(self, x1: int | float | complex | Array, x2: int | float | complex | Array) -> Array: - return Array(np.less(unwrap(x1), unwrap(x2))) + return Array(cp.less(unwrap(x1), unwrap(x2))) def less_equal(self, x1: int | float | complex | Array, x2: int | float | complex | Array) -> Array: - return Array(np.less_equal(unwrap(x1), unwrap(x2))) + return Array(cp.less_equal(unwrap(x1), unwrap(x2))) def greater(self, x1: int | float | complex | Array, x2: int | float | complex | Array) -> Array: - return Array(np.greater(unwrap(x1), unwrap(x2))) + return Array(cp.greater(unwrap(x1), unwrap(x2))) def greater_equal(self, x1: int | float | complex | Array, x2: int | float | complex | Array) -> Array: - return Array(np.greater_equal(unwrap(x1), unwrap(x2))) + return Array(cp.greater_equal(unwrap(x1), unwrap(x2))) # Bitwise def bitwise_and(self, x1: bool | int | Array, x2: bool | int | Array) -> Array: - return Array(np.bitwise_and(unwrap(x1), unwrap(x2))) + return Array(cp.bitwise_and(unwrap(x1), unwrap(x2))) def iand[T: Array](self, x1: T, x2: bool | int | Array) -> T: x1.value &= unwrap(x2) return x1 def bitwise_invert(self, x: Array) -> Array: - return Array(np.bitwise_not(x.value)) + return Array(cp.bitwise_invert(x.value)) def bitwise_or(self, x1: bool | int | Array, x2: bool | int | Array) -> Array: - return Array(np.bitwise_or(unwrap(x1), unwrap(x2))) + return Array(cp.bitwise_or(unwrap(x1), unwrap(x2))) def ior[T: Array](self, x1: T, x2: bool | int | Array) -> T: x1.value |= unwrap(x2) return x1 def bitwise_xor(self, x1: bool | int | Array, x2: bool | int | Array) -> Array: - return Array(np.bitwise_xor(unwrap(x1), unwrap(x2))) + return Array(cp.bitwise_xor(unwrap(x1), unwrap(x2))) def ixor[T: Array](self, x1: T, x2: bool | int | Array) -> T: x1.value ^= unwrap(x2) return x1 def bitwise_left_shift(self, x1: int | Array, x2: int | Array) -> Array: - return Array(np.left_shift(unwrap(x1), unwrap(x2))) + return Array(cp.bitwise_left_shift(unwrap(x1), unwrap(x2))) def ilshift[T: Array](self, x1: T, x2: int | Array) -> T: x1.value <<= unwrap(x2) return x1 def bitwise_right_shift(self, x1: int | Array, x2: int | Array) -> Array: - return Array(np.right_shift(unwrap(x1), unwrap(x2))) + return Array(cp.bitwise_right_shift(unwrap(x1), unwrap(x2))) def irshift[T: Array](self, x1: T, x2: int | Array) -> T: x1.value >>= unwrap(x2) @@ -330,22 +317,16 @@ def irshift[T: Array](self, x1: T, x2: int | Array) -> T: # Operators def sign(self, x: Array) -> Array: - return Array(np.sign(x.value)) + return Array(cp.sign(x.value)) def maximum(self, x1: int | float | complex | Array, x2: int | float | complex | Array) -> Array: - return Array(np.maximum(unwrap(x1), unwrap(x2))) + return Array(cp.maximum(unwrap(x1), unwrap(x2))) def argmax(self, x: Array, axis: int | None = None, keepdims: bool = False) -> Array: - v = cast("np.ndarray[Any, Any]", x.value) - if keepdims: - return Array(np.argmax(v, axis=axis, keepdims=True)) - return Array(np.argmax(v, axis=axis)) + return Array(cp.argmax(x.value, axis=axis, keepdims=keepdims)) def argmin(self, x: Array, axis: int | None = None, keepdims: bool = False) -> Array: - v = cast("np.ndarray[Any, Any]", x.value) - if keepdims: - return Array(np.argmin(v, axis=axis, keepdims=True)) - return Array(np.argmin(v, axis=axis)) + return Array(cp.argmin(x.value, axis=axis, keepdims=keepdims)) def set_item(self, x: Array, key: ArrayKey, value: bool | int | float | complex | Array) -> None: x.value[key] = unwrap(value) @@ -354,29 +335,23 @@ def get_item(self, x: Array, key: ArrayKey) -> Array: return Array(x.value[key]) # RNG + # CuPy is still limited in how random states are handled; cp.random.Generator does not expose a serializable state; + # also, the function cupy.random.get_random_state() returns a cp.random.RandomState for the current device, but this + # object is not serializable. To work around this limitation, CupyBackend has a _seed attribute that stores the seed + # at any time, and this seed is returned and set by get_rng_state and set_rng_state, respectively. + # This might result in some loss of reproducibility, but it allows overall compatibility with the codebase. def set_seed(self, seed: int) -> None: - # Seed both the legacy global state and our owned Generator. The legacy state is - # important because some downstream libraries (sklearn, pandas) consult it. - np.random.seed(seed) # noqa: NPY002 - self._rng = np.random.default_rng(seed) + self._seed = seed + self._rng = cp.random.default_rng(self._seed) def get_rng_state(self) -> dict[str, Any]: - # ``np.random.get_state()`` returns a tuple by default; ``legacy=False`` returns - # the equivalent dict form, which both matches the surrounding ``dict[str, Any]`` - # value type (so mypyc's strict union narrowing is satisfied) and round-trips - # cleanly through ``np.random.set_state``. - return { - "numpy_bit_generator_state": deepcopy(self._rng.bit_generator.state), - "numpy_legacy_state": np.random.get_state(legacy=False), # noqa: NPY002 - } + return {"cupy_generator_seed": self._seed} def set_rng_state(self, state: dict[str, Any]) -> None: - if "numpy_bit_generator_state" in state: - self._rng = np.random.default_rng() - self._rng.bit_generator.state = state["numpy_bit_generator_state"] - if "numpy_legacy_state" in state: - np.random.set_state(state["numpy_legacy_state"]) # noqa: NPY002 + if "cupy_generator_seed" in state: + self._seed = state["cupy_generator_seed"] + self._rng = cp.random.default_rng(self._seed) def normal(self, mean: float = 0.0, std: float = 1.0, shape: tuple[int, ...] = ()) -> Array: return Array(self._rng.normal(loc=mean, scale=std, size=shape)) @@ -385,10 +360,10 @@ def uniform(self, low: float = 0.0, high: float = 1.0, shape: tuple[int, ...] = return Array(self._rng.uniform(low=low, high=high, size=shape)) def normal_like(self, x: Array, mean: float = 0.0, std: float = 1.0) -> Array: - return Array(self._rng.normal(loc=mean, scale=std, size=x.value.shape).astype(x.value.dtype)) + return Array(self._rng.normal(loc=mean, scale=std, size=x.value.shape, dtype=x.dtype.backend_dtype)) def uniform_like(self, x: Array, low: float = 0.0, high: float = 1.0) -> Array: - return Array(self._rng.uniform(low=low, high=high, size=x.value.shape).astype(x.value.dtype)) + return Array(self._rng.uniform(low=low, high=high, size=x.value.shape, dtype=x.dtype.backend_dtype)) def choice(self, x: Array, size: int, replace: bool = True) -> Array: return Array(self._rng.choice(x.value, size=size, replace=replace)) diff --git a/decent_array/interoperability/_jax/jax_backend.py b/decent_array/interoperability/_jax/jax_backend.py index bdc40eb..6c7cec8 100644 --- a/decent_array/interoperability/_jax/jax_backend.py +++ b/decent_array/interoperability/_jax/jax_backend.py @@ -23,6 +23,7 @@ from decent_array._errors import ( MatrixTransposeError, NDimError, + NotScalarError, UnsupportedDeviceError, UnsupportedDTypeCreationError, stack_empty_error, @@ -66,7 +67,7 @@ def device_to_native(self, device: Devices) -> jax.Device: return jax.devices("cpu")[0] if device == Devices.GPU: return jax.devices("gpu")[0] - raise ValueError(f"Unsupported device for JAX: {device}") + raise UnsupportedDeviceError(self.name, device.value) def device_of(self, x: Array) -> Devices: platform = x.value.device.platform @@ -103,7 +104,7 @@ def to_scalar(self, x: Array) -> Any: # noqa: ANN401 """ if not is_scalar(x): - raise TypeError("Only 0-dim arrays can be converted to Python scalars.") + raise NotScalarError(x.ndim) return x.value.item() def stack(self, arrays: Sequence[Array], axis: int = 0) -> Array: diff --git a/decent_array/interoperability/_numpy/numpy_backend.py b/decent_array/interoperability/_numpy/numpy_backend.py index 9a6577b..555c4b9 100644 --- a/decent_array/interoperability/_numpy/numpy_backend.py +++ b/decent_array/interoperability/_numpy/numpy_backend.py @@ -18,6 +18,7 @@ from decent_array._errors import ( MatrixTransposeError, NDimError, + NotScalarError, UnsupportedDeviceError, UnsupportedDTypeCreationError, stack_empty_error, @@ -96,7 +97,7 @@ def to_scalar(self, x: Array) -> Any: # noqa: ANN401 """ if not is_scalar(x): - raise TypeError("Only 0-dim arrays can be converted to Python scalars.") + raise NotScalarError(x.ndim) return x.value.item() def stack(self, arrays: Sequence[Array], axis: int = 0) -> Array: diff --git a/decent_array/interoperability/_pytorch/pytorch_backend.py b/decent_array/interoperability/_pytorch/pytorch_backend.py index 542a534..b6df722 100644 --- a/decent_array/interoperability/_pytorch/pytorch_backend.py +++ b/decent_array/interoperability/_pytorch/pytorch_backend.py @@ -15,7 +15,14 @@ from numpy.typing import NDArray from decent_array import Array -from decent_array._errors import MatrixTransposeError, NDimError, UnsupportedDTypeCreationError, stack_empty_error +from decent_array._errors import ( + MatrixTransposeError, + NDimError, + NotScalarError, + UnsupportedDeviceError, + UnsupportedDTypeCreationError, + stack_empty_error, +) from decent_array._utils import is_scalar, unwrap from decent_array.interoperability._abstracts import Backend from decent_array.interoperability._backend_manager import register_backend @@ -55,7 +62,7 @@ def device_to_native(self, device: Devices) -> str: return "cuda" if device == Devices.MPS: return "mps" - raise ValueError(f"Unsupported device: {device}") + raise UnsupportedDeviceError(self.name, device.value) def device_of(self, x: Array) -> Devices: kind = x.value.device.type @@ -100,7 +107,7 @@ def to_scalar(self, x: Array) -> Any: # noqa: ANN401 """ if not is_scalar(x): - raise TypeError("Only 0-dim arrays can be converted to Python scalars.") + raise NotScalarError(x.ndim) return x.value.item() def stack(self, arrays: Sequence[Array], axis: int = 0) -> Array: diff --git a/decent_array/interoperability/_tensorflow/tensorflow_backend.py b/decent_array/interoperability/_tensorflow/tensorflow_backend.py index 0bd3f51..2742784 100644 --- a/decent_array/interoperability/_tensorflow/tensorflow_backend.py +++ b/decent_array/interoperability/_tensorflow/tensorflow_backend.py @@ -21,6 +21,7 @@ from decent_array._errors import ( MatrixTransposeError, NDimError, + NotScalarError, UnsupportedDeviceError, UnsupportedDTypeCreationError, stack_empty_error, @@ -65,7 +66,7 @@ def eye(self, n: int) -> Array: def device_to_native(self, device: Devices) -> str: if device in {Devices.CPU, Devices.GPU}: return f"/{device.value}:0" - raise ValueError(f"Unsupported device for TensorFlow: {device}") + raise UnsupportedDeviceError(self.name, device.value) def device_of(self, x: Array) -> Devices: device_str = x.value.device.lower() @@ -113,7 +114,7 @@ def to_scalar(self, x: Array) -> Any: # noqa: ANN401 """ if not is_scalar(x): - raise TypeError("Only 0-dim arrays can be converted to Python scalars.") + raise NotScalarError(x.ndim) return x.value.numpy().item() def stack(self, arrays: Sequence[Array], axis: int = 0) -> Array: diff --git a/docs/source/user.rst b/docs/source/user.rst index 5f9c298..798c474 100644 --- a/docs/source/user.rst +++ b/docs/source/user.rst @@ -220,6 +220,7 @@ reliable way is to try operations that involve the dtype and observe if the fram - ✓ - ✓ - ✓ + - ✓ - * - ``bfloat16`` - ✗ From a9644bc55296bca0380d2ba10e9f5c4e7cfc64db Mon Sep 17 00:00:00 2001 From: nicola-bastianello Date: Thu, 9 Jul 2026 14:12:30 +0200 Subject: [PATCH 23/23] fix: small fixes --- decent_array/interoperability/_cupy/cupy_backend.py | 9 +++++---- docs/source/user.rst | 5 +++++ 2 files changed, 10 insertions(+), 4 deletions(-) diff --git a/decent_array/interoperability/_cupy/cupy_backend.py b/decent_array/interoperability/_cupy/cupy_backend.py index 1939924..0762049 100644 --- a/decent_array/interoperability/_cupy/cupy_backend.py +++ b/decent_array/interoperability/_cupy/cupy_backend.py @@ -12,6 +12,7 @@ from typing import Any import cupy as cp +import numpy as np from numpy.typing import NDArray from decent_array import Array @@ -435,22 +436,22 @@ def complex128(self) -> Any: # noqa: ANN401 @property def e(self) -> Any: # noqa: ANN401 """e = 2.71828...""" # noqa: D403 - return cp.e + return np.e @property def inf(self) -> Any: # noqa: ANN401 """Infinity.""" - return cp.inf + return np.inf @property def nan(self) -> Any: # noqa: ANN401 """Not-a-number.""" - return cp.nan + return np.nan @property def pi(self) -> Any: # noqa: ANN401 """pi = 3.14159...""" # noqa: D403 - return cp.pi + return np.pi register_backend(Frameworks.CUPY, CupyBackend) diff --git a/docs/source/user.rst b/docs/source/user.rst index 798c474..a6aa796 100644 --- a/docs/source/user.rst +++ b/docs/source/user.rst @@ -181,6 +181,7 @@ reliable way is to try operations that involve the dtype and observe if the fram - - - + - * - ``uint8`` - ✓ - ✓ @@ -215,6 +216,7 @@ reliable way is to try operations that involve the dtype and observe if the fram - - - + - * - ``float16`` - ✓ - ✓ @@ -256,6 +258,7 @@ reliable way is to try operations that involve the dtype and observe if the fram - - - + - * - ``complex64`` - ✓ - ✓ @@ -283,6 +286,7 @@ reliable way is to try operations that involve the dtype and observe if the fram - - - + - * - ``qint8`` - ✗ - ✗ @@ -324,6 +328,7 @@ reliable way is to try operations that involve the dtype and observe if the fram - - - + - * - ``unicode_`` - ✓ - ✗