New function at()

.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"]  # Do not release
 
 [project.optional-dependencies]
 tests = [
@@ -63,9 +63,11 @@ platforms = ["linux-64", "osx-arm64", "win-64"]
 
 [tool.pixi.dependencies]
 python = ">=3.10.15,<3.14"
-array-api-compat = ">=1.1.1"
+# array-api-compat = ">=1.10.0"  # Do not release
 
 [tool.pixi.pypi-dependencies]
+# Do not release: main at least @ gh#205
+array-api-compat = { git = "https://github.com/data-apis/array-api-compat.git" }
 array-api-extra = { path = ".", editable = true }
 
 [tool.pixi.feature.lint.dependencies]
@@ -99,7 +101,9 @@ tests-cov = "pytest -v -ra --cov --cov-report=xml --cov-report=term --durations=
 
 clean-vendor-compat = "rm -rf vendor_tests/array_api_compat"
 clean-vendor-extra = "rm -rf vendor_tests/array_api_extra"
-copy-vendor-compat = { cmd = "cp -r $(python -c 'import site; print(site.getsitepackages()[0])')/array_api_compat vendor_tests/", depends-on = ["clean-vendor-compat"] }
+copy-vendor-compat = { cmd = "cp -r $(python -c 'import site; print(site.getsitepackages()[0])')/array_api_compat vendor_tests/", depends-on = [
+    "clean-vendor-compat",
+] }
 copy-vendor-extra = { cmd = "cp -r src/array_api_extra vendor_tests/", depends-on = ["clean-vendor-extra"] }
 tests-vendor = { cmd = "pytest -v vendor_tests", depends-on = ["copy-vendor-compat", "copy-vendor-extra"] }
 
@@ -130,6 +134,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 +171,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 +230,8 @@ reportAny = false
 reportExplicitAny = false
 # data-apis/array-api-strict#6
 reportUnknownMemberType = false
+# no array-api-compat type stubs
+reportUnknownVariableType = false
 
 
 # Ruff
@@ -236,6 +273,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 +309,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,12 +1,22 @@
 """Extra array functions built on top of the array API standard."""
 
-from ._funcs import atleast_nd, cov, create_diagonal, expand_dims, kron, setdiff1d, sinc
+from ._funcs import (
+    at,
+    atleast_nd,
+    cov,
+    create_diagonal,
+    expand_dims,
+    kron,
+    setdiff1d,
+    sinc,
+)
 
 __version__ = "0.4.1.dev0"
 
 # pylint: disable=duplicate-code
 __all__ = [
     "__version__",
+    "at",
     "atleast_nd",
     "cov",
     "create_diagonal",