Array API compatibility for pytorch and jax

.github/workflows/test_and_deploy.yml

@@ -24,7 +24,7 @@ jobs:
     strategy:
       matrix:
         platform: [ubuntu-latest, windows-latest, macos-latest]
-        python-version: ['3.9', '3.10', '3.11', '3.12', '3.13']
+        python-version: ['3.9', '3.10', '3.11', '3.12', '3.13', '3.14']
 
     steps:
       - uses: actions/checkout@v4
@@ -43,11 +43,18 @@ jobs:
           python -m pip install --upgrade pip
           python -m pip install setuptools tox tox-gh-actions
 
+      - name: Set up MSVC
+        if: runner.os == 'Windows'
+        uses: TheMrMilchmann/setup-msvc-dev@v4
+        with:
+          arch: x64
+
       # this runs the platform-specific tests declared in tox.ini
       - name: Test with tox
-        run: python -m tox
         env:
           PLATFORM: ${{ matrix.platform }}
+          TORCHINDUCTOR_CACHE_DIR: ${{ runner.temp }}/torchinductor
+        run: python -m tox
 
       - name: Upload coverage reports to Codecov
         uses: codecov/codecov-action@v4

examples/skip_performance_basis_functions.py

@@ -0,0 +1,54 @@
+import collections
+import time
+
+import jax.numpy as jax
+import matplotlib.pyplot as plt
+import numpy as np
+import pandas as pd
+import seaborn as sns
+import splinebox
+import torch
+
+results = collections.defaultdict(list)
+
+for length in [10, 100, 1000, 10000, 100000]:
+    for replicate in np.arange(11):
+        x = np.random.rand(length)
+        for array_namespace in ["numpy", "torch", "jax"]:
+            if array_namespace == "torch":
+                x = torch.tensor(x)
+            elif array_namespace == "jax":
+                x = jax.asarray(x)
+            for basis_function_name, basis_function_class in splinebox.basis_functions.inventory().items():
+                if "Exponential" in basis_function_name:
+                    basis_function = basis_function_class(5)
+                else:
+                    basis_function = basis_function_class()
+                for deriv in range(3):
+                    for jit in [True, False]:
+                        start = time.perf_counter_ns()
+                        basis_function(x, derivative=deriv, jit=jit)
+                        stop = time.perf_counter_ns()
+
+                        results["length"].append(length)
+                        results["Replicate"].append(replicate)
+                        # results["Array namespace"].append(array_namespace)
+                        results["Basis function"].append(basis_function_name)
+                        results["Derivative"].append(deriv)
+                        results["JIT"].append(jit)
+                        results["time [ns]"].append(stop - start)
+
+df = pd.DataFrame(results)
+
+g = sns.relplot(
+    data=df.loc[df["Replicate"] > 0],
+    x="length",
+    y="time [ns]",
+    col="Basis function",
+    hue="JIT",
+    row="Array namespace",
+    kind="line",
+)
+g.set(xscale="log")
+g.set(yscale="log")
+plt.show()

pyproject.toml

@@ -94,6 +94,9 @@ test = [
     "tox",
     "pytest",
     "pytest-cov",
+    "array-api-compat",
+    "torch",
+    "jax"
 ]
 examples = [
     "matplotlib",