ENH: implement at (#53)

* ENH: add new function `at`

* MAINT: use released array-api-compat

* update lock-file

* Update dependencies

* Add xpx namespace in documentation

* Change copy to default to None

* raise on incompatible cast

* Update tests/test_at.py

---------

Co-authored-by: Lucas Colley <[email protected]>

.github/workflows/ci.yml

@@ -48,7 +48,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        environment: [ci-py310, ci-py313]
+        environment: [ci-py310, ci-py313, ci-backends]
         runs-on: [ubuntu-latest]
 
     steps:

docs/api-reference.md

@@ -6,6 +6,7 @@
     :nosignatures:
     :toctree: generated
 
+    at
     atleast_nd
     cov
     create_diagonal

docs/conf.py

@@ -53,6 +53,7 @@
 
 intersphinx_mapping = {
     "python": ("https://docs.python.org/3", None),
+    "jax": ("https://jax.readthedocs.io/en/latest", None),
 }
 
 nitpick_ignore = [

pyproject.toml

@@ -26,7 +26,7 @@ classifiers = [
     "Typing :: Typed",
 ]
 dynamic = ["version"]
-dependencies = ["array-api-compat>=1.1.1"]
+dependencies = ["array-api-compat>=1.10.0,<2"]
 
 [project.optional-dependencies]
 tests = [
@@ -62,8 +62,8 @@ channels = ["https://prefix.dev/conda-forge"]
 platforms = ["linux-64", "osx-arm64", "win-64"]
 
 [tool.pixi.dependencies]
-python = ">=3.10.15,<3.14"
-array-api-compat = ">=1.1.1"
+python = ">=3.10,<3.14"
+array-api-compat = ">=1.10.0,<2"
 
 [tool.pixi.pypi-dependencies]
 array-api-extra = { path = ".", editable = true }
@@ -130,6 +130,35 @@ python = "~=3.10.0"
 [tool.pixi.feature.py313.dependencies]
 python = "~=3.13.0"
 
+# Backends that can run on CPU-only hosts
+[tool.pixi.feature.backends.target.linux-64.dependencies]
+pytorch = "*"
+dask = "*"
+sparse = ">=0.15"
+jax = "*"
+
+[tool.pixi.feature.backends.target.osx-arm64.dependencies]
+pytorch = "*"
+dask = "*"
+sparse = ">=0.15"
+jax = "*"
+
+[tool.pixi.feature.backends.target.win-64.dependencies]
+# pytorch = "*"  # Package unavailable on Windows
+dask = "*"
+sparse = ">=0.15"
+# jax = "*"  # Package unavailable on Windows
+
+# Backends that require a GPU host and a CUDA driver
+[tool.pixi.feature.cuda-backends.target.linux-64.dependencies]
+cupy = "*"
+
+[tool.pixi.feature.cuda-backends.target.osx-arm64.dependencies]
+# cupy = "*"  # Package unavailable on macOSX
+
+[tool.pixi.feature.cuda-backends.target.win-64.dependencies]
+cupy = "*"
+
 [tool.pixi.environments]
 default = { solve-group = "default" }
 lint = { features = ["lint"], solve-group = "default" }
@@ -138,7 +167,9 @@ docs = { features = ["docs"], solve-group = "default" }
 dev = { features = ["lint", "tests", "docs", "dev"], solve-group = "default" }
 ci-py310 = ["py310", "tests"]
 ci-py313 = ["py313", "tests"]
-
+# CUDA not available on free github actions
+ci-backends = ["py310", "tests", "backends"]
+tests-backends = ["py310", "tests", "backends", "cuda-backends"]
 
 # pytest
 
@@ -195,6 +226,8 @@ reportAny = false
 reportExplicitAny = false
 # data-apis/array-api-strict#6
 reportUnknownMemberType = false
+# no array-api-compat type stubs
+reportUnknownVariableType = false
 
 
 # Ruff
@@ -236,6 +269,7 @@ ignore = [
     "PLR09",   # Too many <...>
     "PLR2004", # Magic value used in comparison
     "ISC001",  # Conflicts with formatter
+    "N801",    # Class name should use CapWords convention
     "N802",    # Function name should be lowercase
     "N806",    # Variable in function should be lowercase
 ]
@@ -271,6 +305,7 @@ checks = [
     "ES01",
 ]
 exclude = [ # don't report on objects that match any of these regex
+    '.*test_at.*',
     '.*test_funcs.*',
     '.*test_utils.*',
     '.*test_version.*',

src/array_api_extra/__init__.py

@@ -1,6 +1,7 @@
 """Extra array functions built on top of the array API standard."""
 
 from ._funcs import (
+    at,
     atleast_nd,
     cov,
     create_diagonal,
@@ -16,6 +17,7 @@
 # pylint: disable=duplicate-code
 __all__ = [
     "__version__",
+    "at",
     "atleast_nd",
     "cov",
     "create_diagonal",