ENH: implement 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,<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",