ndlike 0.1.0__tar.gz

ndlike-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 the probly team
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

ndlike-0.1.0/PKG-INFO

@@ -0,0 +1,160 @@
+Metadata-Version: 2.4
+Name: ndlike
+Version: 0.1.0
+Summary: A toolkit for defining custom array types on top of NumPy, JAX, and PyTorch.
+Author-email: Clemens Damke <clemens@cortys.de>
+License-Expression: MIT
+Project-URL: source, https://github.com/Cortys/ndlike
+Project-URL: tracker, https://github.com/Cortys/ndlike/issues
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Science/Research
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Programming Language :: Python :: 3.14
+Classifier: Programming Language :: Python :: 3.15
+Requires-Python: >=3.12
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: numpy>=2.1.0
+Requires-Dist: flextype>=0.2.0
+Provides-Extra: torch
+Requires-Dist: torch>=2.6.0; extra == "torch"
+Provides-Extra: jax
+Requires-Dist: jax>=0.8.1; extra == "jax"
+Dynamic: license-file
+
+# ndlike
+
+NumPy arrays, JAX arrays, and PyTorch tensors provide similar primitives for
+indexing, reshaping, stacking, and numerical operations. `ndlike` provides tools
+to build wrappers around these types that behave like the underlying arrays,
+but whose elements represent custom composite values—such as vectors or
+distributions—instead of individual numbers.
+
+Protocols and reusable base classes handle array interfaces and backend function
+dispatch. Protected axes distinguish collection dimensions from each element's
+internal dimensions, so indexing, reshaping, or stacking a collection preserves
+the structure of its elements.
+
+## Installation
+
+Requires Python 3.12 or newer.
+
+```bash
+pip install ndlike
+pip install 'ndlike[torch]'
+pip install 'ndlike[jax]'
+```
+
+NumPy and flextype are core dependencies. Torch and JAX are optional and are
+never imported by the backend-independent modules. Import backend modules explicitly.
+
+## Array protocols and base classes
+
+`ndlike.array_like.ArrayLike` describes indexing, shape, and metadata rather than
+just convertibility to an array (as in `numpy.typing.ArrayLike`). It does not
+require the full numerical Array API. `ArrayLikeWithoutSize` covers the shared
+interface without NumPy/JAX's integer `size` property, since Torch uses a method.
+
+Backend-specific APIs build on these protocols:
+
+| Module | Protocol | Conversion helper | Conversion hook |
+| --- | --- | --- | --- |
+| `ndlike.array_like` | `NumpyLike` | `to_numpy_like` | `__array_like__` |
+| `ndlike.torch_like` | `TorchLike` | `to_torch_like` | `__torch_like__` |
+| `ndlike.jax_like` | `JaxLike` | `to_jax_like` | `__jax_like__` |
+
+Each module also provides matching `Implementation` and `Convertible` classes,
+such as `NumpyLikeImplementation` and `NumpyLikeConvertible`.
+
+Use the array protocols to type inputs, or subclass an implementation base to
+reuse convenience methods while supplying your representation's metadata and
+operation-dispatch hooks. NumPy uses `__array_function__` / `__array_ufunc__`,
+Torch uses `__torch_function__`, and ndlike supplies `__jax_function__` for JAX.
+
+Conversion helpers call the conversion hook when an input satisfies the matching
+convertible protocol: the core array interface plus the hook, not the hook alone.
+This lets custom types preserve their structure instead of becoming raw arrays.
+Otherwise, the helpers fall back to `numpy.asanyarray`, `torch.as_tensor`, and
+`jax.numpy.asarray`, respectively. NumPy retains `__array_like__` for backward
+compatibility and consistency with NumPy's `__array__`.
+
+## Protected axes
+
+Protected-axis base classes give structured objects array-like behavior over
+their batch dimensions while preserving trailing internal dimensions. For
+example, a batch of six 3D vectors should have logical shape `(6,)`, not `(6, 3)`:
+
+```python
+from dataclasses import dataclass
+from typing import ClassVar
+
+import numpy as np
+
+from ndlike.protected_axis.numpy import NumpyAxisProtected
+
+
+@dataclass
+class Vectors(NumpyAxisProtected[np.ndarray]):
+    array: np.ndarray
+    protected_axes: ClassVar[dict[str, int]] = {"array": 1}
+
+
+vectors = Vectors(np.ones((6, 3)))
+assert vectors.shape == (6,)
+assert vectors.protected_shape == (3,)
+assert vectors[0].array.shape == (3,)
+
+grid = vectors.reshape(2, 3)
+assert grid.shape == (2, 3)
+assert grid.array.shape == (2, 3, 3)
+
+stacked = np.stack([vectors, vectors])
+assert isinstance(stacked, Vectors)
+assert stacked.shape == (2, 6)
+```
+
+`protected_axes` maps field names to the number of trailing axes to preserve.
+Multiple fields can have different internal shapes while sharing the same batch
+shape; indexing and supported shape operations update them together. Dataclasses
+work out of the box; other classes can override `with_protected_values` to
+reconstruct results.
+
+Numerical operations are opt-in: use `permitted_functions` (for example,
+`{np.mean}`) and, for NumPy ufuncs, `permitted_ufuncs` (for example,
+`{np.add: ["__call__", "reduce"]}`). Supported reductions operate on batch axes,
+leaving internal axes intact.
+
+The corresponding backend classes are `TorchAxisProtected` in
+`ndlike.protected_axis.torch` and `JaxAxisProtected` in
+`ndlike.protected_axis.jax`.
+
+## JAX function dispatch
+
+JAX has no native equivalent of NumPy's `__array_function__`. For custom objects,
+use wrappers from `ndlike.jax_functions`, such as `jax_reshape`, `jax_stack`, and
+`jax_mean`, instead of the corresponding `jax.numpy` functions. They dispatch to
+`__jax_function__` on custom types and call JAX normally for native arrays.
+`jax_function_dispatch` lets you add the same dispatch behavior to your own
+functions. `JaxLikeImplementation` also provides dataclass-based pytree support
+for JAX transformations such as `jax.jit` and `jax.vmap`.
+
+## Development
+
+```bash
+uv sync --all-extras
+uv run pytest
+uv run prek install
+uv run prek run --all-files
+uv build
+```
+
+Dependency groups follow flextype: `lint` contains Ruff and ty, `test` contains
+pytest, and `dev` includes both groups plus prek. Development Torch wheels use
+the CPU index on Linux and Windows. Type checking requires both backend extras.
+
+`scripts/check_backend_naming.py` enforces backend prefixes on definitions
+(BKN001) and restricts runtime Torch/JAX imports to their respective
+backend-prefixed modules (BKN002). Type-only imports are exempt. Intentional
+exceptions use an explicit `# noqa: BKN001` or `# noqa: BKN002` on the opening line.

ndlike-0.1.0/README.md

@@ -0,0 +1,134 @@
+# ndlike
+
+NumPy arrays, JAX arrays, and PyTorch tensors provide similar primitives for
+indexing, reshaping, stacking, and numerical operations. `ndlike` provides tools
+to build wrappers around these types that behave like the underlying arrays,
+but whose elements represent custom composite values—such as vectors or
+distributions—instead of individual numbers.
+
+Protocols and reusable base classes handle array interfaces and backend function
+dispatch. Protected axes distinguish collection dimensions from each element's
+internal dimensions, so indexing, reshaping, or stacking a collection preserves
+the structure of its elements.
+
+## Installation
+
+Requires Python 3.12 or newer.
+
+```bash
+pip install ndlike
+pip install 'ndlike[torch]'
+pip install 'ndlike[jax]'
+```
+
+NumPy and flextype are core dependencies. Torch and JAX are optional and are
+never imported by the backend-independent modules. Import backend modules explicitly.
+
+## Array protocols and base classes
+
+`ndlike.array_like.ArrayLike` describes indexing, shape, and metadata rather than
+just convertibility to an array (as in `numpy.typing.ArrayLike`). It does not
+require the full numerical Array API. `ArrayLikeWithoutSize` covers the shared
+interface without NumPy/JAX's integer `size` property, since Torch uses a method.
+
+Backend-specific APIs build on these protocols:
+
+| Module | Protocol | Conversion helper | Conversion hook |
+| --- | --- | --- | --- |
+| `ndlike.array_like` | `NumpyLike` | `to_numpy_like` | `__array_like__` |
+| `ndlike.torch_like` | `TorchLike` | `to_torch_like` | `__torch_like__` |
+| `ndlike.jax_like` | `JaxLike` | `to_jax_like` | `__jax_like__` |
+
+Each module also provides matching `Implementation` and `Convertible` classes,
+such as `NumpyLikeImplementation` and `NumpyLikeConvertible`.
+
+Use the array protocols to type inputs, or subclass an implementation base to
+reuse convenience methods while supplying your representation's metadata and
+operation-dispatch hooks. NumPy uses `__array_function__` / `__array_ufunc__`,
+Torch uses `__torch_function__`, and ndlike supplies `__jax_function__` for JAX.
+
+Conversion helpers call the conversion hook when an input satisfies the matching
+convertible protocol: the core array interface plus the hook, not the hook alone.
+This lets custom types preserve their structure instead of becoming raw arrays.
+Otherwise, the helpers fall back to `numpy.asanyarray`, `torch.as_tensor`, and
+`jax.numpy.asarray`, respectively. NumPy retains `__array_like__` for backward
+compatibility and consistency with NumPy's `__array__`.
+
+## Protected axes
+
+Protected-axis base classes give structured objects array-like behavior over
+their batch dimensions while preserving trailing internal dimensions. For
+example, a batch of six 3D vectors should have logical shape `(6,)`, not `(6, 3)`:
+
+```python
+from dataclasses import dataclass
+from typing import ClassVar
+
+import numpy as np
+
+from ndlike.protected_axis.numpy import NumpyAxisProtected
+
+
+@dataclass
+class Vectors(NumpyAxisProtected[np.ndarray]):
+    array: np.ndarray
+    protected_axes: ClassVar[dict[str, int]] = {"array": 1}
+
+
+vectors = Vectors(np.ones((6, 3)))
+assert vectors.shape == (6,)
+assert vectors.protected_shape == (3,)
+assert vectors[0].array.shape == (3,)
+
+grid = vectors.reshape(2, 3)
+assert grid.shape == (2, 3)
+assert grid.array.shape == (2, 3, 3)
+
+stacked = np.stack([vectors, vectors])
+assert isinstance(stacked, Vectors)
+assert stacked.shape == (2, 6)
+```
+
+`protected_axes` maps field names to the number of trailing axes to preserve.
+Multiple fields can have different internal shapes while sharing the same batch
+shape; indexing and supported shape operations update them together. Dataclasses
+work out of the box; other classes can override `with_protected_values` to
+reconstruct results.
+
+Numerical operations are opt-in: use `permitted_functions` (for example,
+`{np.mean}`) and, for NumPy ufuncs, `permitted_ufuncs` (for example,
+`{np.add: ["__call__", "reduce"]}`). Supported reductions operate on batch axes,
+leaving internal axes intact.
+
+The corresponding backend classes are `TorchAxisProtected` in
+`ndlike.protected_axis.torch` and `JaxAxisProtected` in
+`ndlike.protected_axis.jax`.
+
+## JAX function dispatch
+
+JAX has no native equivalent of NumPy's `__array_function__`. For custom objects,
+use wrappers from `ndlike.jax_functions`, such as `jax_reshape`, `jax_stack`, and
+`jax_mean`, instead of the corresponding `jax.numpy` functions. They dispatch to
+`__jax_function__` on custom types and call JAX normally for native arrays.
+`jax_function_dispatch` lets you add the same dispatch behavior to your own
+functions. `JaxLikeImplementation` also provides dataclass-based pytree support
+for JAX transformations such as `jax.jit` and `jax.vmap`.
+
+## Development
+
+```bash
+uv sync --all-extras
+uv run pytest
+uv run prek install
+uv run prek run --all-files
+uv build
+```
+
+Dependency groups follow flextype: `lint` contains Ruff and ty, `test` contains
+pytest, and `dev` includes both groups plus prek. Development Torch wheels use
+the CPU index on Linux and Windows. Type checking requires both backend extras.
+
+`scripts/check_backend_naming.py` enforces backend prefixes on definitions
+(BKN001) and restricts runtime Torch/JAX imports to their respective
+backend-prefixed modules (BKN002). Type-only imports are exempt. Intentional
+exceptions use an explicit `# noqa: BKN001` or `# noqa: BKN002` on the opening line.

ndlike-0.1.0/pyproject.toml

@@ -0,0 +1,98 @@
+[build-system]
+requires = ["setuptools>=77", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "ndlike"
+version = "0.1.0"
+description = "A toolkit for defining custom array types on top of NumPy, JAX, and PyTorch."
+readme = "README.md"
+license = "MIT"
+requires-python = ">=3.12"
+dependencies = ["numpy>=2.1.0", "flextype>=0.2.0"]
+authors = [{ name = "Clemens Damke", email = "clemens@cortys.de" }]
+classifiers = [
+    "Development Status :: 3 - Alpha",
+    "Intended Audience :: Developers",
+    "Intended Audience :: Science/Research",
+    "Programming Language :: Python :: 3.12",
+    "Programming Language :: Python :: 3.13",
+    "Programming Language :: Python :: 3.14",
+    "Programming Language :: Python :: 3.15",
+]
+
+[project.optional-dependencies]
+torch = ["torch>=2.6.0"]
+jax = ["jax>=0.8.1"]
+
+[project.urls]
+source = "https://github.com/Cortys/ndlike"
+tracker = "https://github.com/Cortys/ndlike/issues"
+
+[tool.setuptools.packages.find]
+where = ["src"]
+include = ["ndlike", "ndlike.*"]
+
+[tool.setuptools.package-data]
+ndlike = ["py.typed"]
+
+[dependency-groups]
+lint = ["ruff>=0.17.1", "ty>=0.0.89"]
+test = ["pytest>=8.0", "typing-extensions>=4.7.0"]
+dev = [
+    { include-group = "lint" },
+    { include-group = "test" },
+    "prek>=0.5.6",
+]
+
+[tool.ruff]
+line-length = 120
+target-version = "py312"
+output-format = "full"
+
+[tool.ruff.lint]
+select = ["ALL"]
+external = ["BKN001", "BKN002"]
+ignore = ["COM812", "CPY001", "FBT", "FIX002", "N812", "PLR0913", "PLR0917", "PLR2004", "RET504"]
+
+[tool.ruff.lint.pydocstyle]
+convention = "google"
+
+[tool.ruff.lint.per-file-ignores]
+"tests/**/*.py" = [
+    "ANN001", "ANN201", "ANN202", "C901", "D100", "D101", "D102", "D103", "D404",
+    "FBT001", "INP001", "N803", "PD901", "PLR2004", "S101", "TCH",
+]
+"__init__.py" = ["I002", "PLC0414"]
+# Previously private modules retain their original documentation during extraction.
+"src/ndlike/protected_axis/*.py" = ["D102", "D103", "D104", "D105"]
+
+[tool.ruff.lint.mccabe]
+max-complexity = 15
+
+[tool.ruff.lint.isort]
+known-first-party = ["ndlike"]
+force-sort-within-sections = true
+combine-as-imports = true
+force-single-line = false
+no-lines-before = ["future"]
+required-imports = ["from __future__ import annotations"]
+extra-standard-library = ["typing_extensions"]
+
+[tool.pytest.ini_options]
+testpaths = ["tests/ndlike"]
+minversion = "8.0"
+pythonpath = ["src"]
+addopts = ["--import-mode=importlib"]
+
+[tool.ty.src]
+include = ["src/ndlike", "tests/ndlike", "scripts"]
+
+# Use small CPU wheels in development and CI; installed users choose their own index.
+[[tool.uv.index]]
+name = "pytorch-cpu"
+url = "https://download.pytorch.org/whl/cpu"
+explicit = true
+
+[tool.uv.sources]
+torch = [{ index = "pytorch-cpu", marker = "sys_platform != 'darwin'" }]

ndlike-0.1.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

ndlike-0.1.0/src/ndlike/__init__.py

@@ -0,0 +1 @@
+"""Tools for defining custom array types across NumPy, PyTorch, and JAX."""

ndlike-0.1.0/src/ndlike/_switchdispatch.py

@@ -0,0 +1,208 @@
+"""An extensible switch dispatch mechanism.
+
+Akin to the builtin functools.singledispatch mechanism,
+but with an equality-based dispatcher.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Iterator, MutableMapping
+from functools import update_wrapper
+from typing import TYPE_CHECKING, Concatenate, Protocol, overload, override
+
+if TYPE_CHECKING:
+    from collections.abc import Callable, Iterable
+    from types import NotImplementedType
+
+
+class _RegistrationSink[K, V](Protocol):
+    """The contravariant write interface used to constrain registered values."""
+
+    def _register_value(self, key: K, value: V) -> None: ...
+
+
+class _Registration[V]:
+    """Store a value while preserving its exact type and identity."""
+
+    def __init__(self, store: Callable[[V], None]) -> None:
+        self._store = store
+
+    # A bound such as ``W: V`` is not legal: TypeVar bounds cannot be generic.
+    # Contravariant storage constrains W to accepted values without widening it
+    # to V, preserving generic classes, overloads, and callable-object attributes.
+    def __call__[W](self: _Registration[W], value: W) -> W:
+        self._store(value)
+        return value
+
+
+class _FunctionRegistration[F]:
+    """Identity-preserving registration with the function decorator's ``f`` keyword."""
+
+    def __init__(self, store: Callable[[F], None]) -> None:
+        self._store = store
+
+    def __call__[W](self: _FunctionRegistration[W], f: W) -> W:
+        self._store(f)
+        return f
+
+
+class switchdispatch[T, **In, Out]:  # noqa: N801
+    """A switch dispatch decorator.
+
+    Similar to functools.singledispatch, but dispatches based on equality
+    rather than type.
+
+    Example:
+        >>> @switchdispatch
+        >>> def func(x):
+        >>>     return "default"
+        >>>
+        >>> @func.register(1)
+        >>> def _(x):
+        >>>     return "one"
+        >>>
+        >>> @func.register(2)
+        >>> def _(x):
+        >>>     return "two"
+        >>>
+        >>> print(func(1))  # Output: "one"
+        >>> print(func(2))  # Output: "two"
+        >>> print(func(3))  # Output: "default"
+    """
+
+    def __init__(self, func: Callable[Concatenate[T, In], Out]) -> None:
+        """Initialize the switchdispatch with the default function."""
+        self._func = func
+        self._registry: dict[object, Callable[Concatenate[T, In], Out]] = {}
+        update_wrapper(self, func, updated=())
+
+    @overload
+    def register(self, key: T) -> _FunctionRegistration[Callable[Concatenate[T, In], Out]]: ...
+
+    @overload
+    def register[F](self: _RegistrationSink[T, F], key: T, f: F) -> F: ...
+
+    def register(
+        self,
+        key: object,
+        f: Callable[Concatenate[T, In], Out] | None = None,
+    ) -> Callable[Concatenate[T, In], Out] | _FunctionRegistration[Callable[Concatenate[T, In], Out]]:
+        """Register a function for the given key, returning it unchanged."""
+        if f is None:
+            return self.multi_register([key])
+        return self.multi_register([key], f)
+
+    @overload
+    def multi_register(self, keys: Iterable[object]) -> _FunctionRegistration[Callable[Concatenate[T, In], Out]]: ...
+
+    @overload
+    def multi_register[F](self: _RegistrationSink[object, F], keys: Iterable[object], f: F) -> F: ...
+
+    def multi_register(
+        self,
+        keys: Iterable[object],
+        f: Callable[Concatenate[T, In], Out] | None = None,
+    ) -> Callable[Concatenate[T, In], Out] | _FunctionRegistration[Callable[Concatenate[T, In], Out]]:
+        """Register a function for the given keys, returning it unchanged."""
+        if f is not None:
+            for key in keys:
+                self._register_value(key, f)
+            return f
+
+        def store(f: Callable[Concatenate[T, In], Out]) -> None:
+            for key in keys:
+                self._register_value(key, f)
+
+        return _FunctionRegistration(store)
+
+    def _register_value(self, key: object, value: Callable[Concatenate[T, In], Out]) -> None:
+        self._registry[key] = value
+
+    def __call__(self, arg: T, *args: In.args, **kwargs: In.kwargs) -> Out:
+        """Call the appropriate function based on the argument."""
+        func_impl = self._registry.get(arg, self._func)
+        return func_impl(arg, *args, **kwargs)
+
+
+class switch[K, V](MutableMapping[K, V]):  # noqa: N801
+    """A simple switch-case mechanism based on a dictionary."""
+
+    _registry: dict[K, V]
+
+    def __init__(
+        self,
+        cases: dict[K, V] | None = None,
+    ) -> None:
+        """Initialize the switch with the given cases and default value."""
+        self._registry = cases or {}
+
+    @overload
+    def register(self, key: K) -> _Registration[V]: ...
+
+    @overload
+    def register[W](self: _RegistrationSink[K, W], key: K, value: W) -> W: ...
+
+    def register(
+        self,
+        key: K,
+        value: V | NotImplementedType = NotImplemented,
+    ) -> _Registration[V] | V:
+        """Register a value for the given key, returning it unchanged."""
+        if value is NotImplemented:
+            return self.multi_register([key])
+        return self.multi_register([key], value)
+
+    @overload
+    def multi_register(self, keys: Iterable[K]) -> _Registration[V]: ...
+
+    @overload
+    def multi_register[W](self: _RegistrationSink[K, W], keys: Iterable[K], value: W) -> W: ...
+
+    def multi_register(
+        self,
+        keys: Iterable[K],
+        value: V | NotImplementedType = NotImplemented,
+    ) -> _Registration[V] | V:
+        """Register a value for the given keys, returning it unchanged."""
+        if value is not NotImplemented:
+            for key in keys:
+                self._register_value(key, value)
+            return value
+
+        def store(value: V) -> None:
+            for key in keys:
+                self._register_value(key, value)
+
+        return _Registration(store)
+
+    def _register_value(self, key: K, value: V) -> None:
+        self._registry[key] = value
+
+    @override
+    def __getitem__(self, key: K) -> V:
+        """Return the value corresponding to the given key."""
+        return self._registry[key]
+
+    def __call__(self, key: K) -> V:
+        """Return the value corresponding to the given key."""
+        return self._registry[key]
+
+    @override
+    def __setitem__(self, key: K, value: V) -> None:
+        """Set the value for the given key."""
+        self._registry[key] = value
+
+    @override
+    def __delitem__(self, key: K) -> None:
+        """Delete the value for the given key."""
+        del self._registry[key]
+
+    @override
+    def __iter__(self) -> Iterator[K]:
+        """Return an iterator over the keys in the switch cases."""
+        return iter(self._registry)
+
+    @override
+    def __len__(self) -> int:
+        """Return the number of cases in the switch."""
+        return len(self._registry)