From 3589097c572d5026115f7104deaacc5933774875 Mon Sep 17 00:00:00 2001
From: faymanns <florian.aymanns@epfl.ch>
Date: Mon, 28 Sep 2026 14:26:29 +0200
Subject: [PATCH 01/15] First working version of basis functions with
 array_api_compat

---
 examples/skip_performance_basis_functions.py |  54 ++
 pyproject.toml                               |   3 +
 src/splinebox/basis_functions.py             | 906 ++++++++++++++++---
 src/splinebox/spline_curves.py               |   2 +-
 tests/conftest.py                            |  85 +-
 tests/test_basis_functions.py                |  33 +-
 6 files changed, 959 insertions(+), 124 deletions(-)
 create mode 100644 examples/skip_performance_basis_functions.py

diff --git a/examples/skip_performance_basis_functions.py b/examples/skip_performance_basis_functions.py
new file mode 100644
index 0000000..f43f058
--- /dev/null
+++ b/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()
diff --git a/pyproject.toml b/pyproject.toml
index 794f5d8..1da2a06 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -94,6 +94,9 @@ test = [
     "tox",
     "pytest",
     "pytest-cov",
+    "array-api-compat",
+    "pytorch",
+    "jax"
 ]
 examples = [
     "matplotlib",
diff --git a/src/splinebox/basis_functions.py b/src/splinebox/basis_functions.py
index d69a590..df1c981 100644
--- a/src/splinebox/basis_functions.py
+++ b/src/splinebox/basis_functions.py
@@ -10,15 +10,69 @@
 For more information on implementing a new basis function, see :class:`splinebox.basis_functions.BasisFunction`.
 """
 
+import contextlib
 import inspect
 import math
 import sys
 import warnings
 
+import array_api_compat
+import array_api_compat.numpy
 import numba
 import numpy as np
 
 
+@contextlib.contextmanager
+def _torch_compile_recompile_limit(limit=64):
+    """
+    Context manager that temporarily raises the torch dynamo recompile limit.
+
+    Compiling the filters with torch.compile unrolls their loops statically,
+    which means a separate graph is compiled for every combination of input
+    shape and basis function parameters (e.g. the pole of the exponential
+    basis function depends on ``M``). This can exceed the default recompile
+    limit of a single code object, so it is temporarily raised while compiling.
+    """
+    import torch
+
+    config = torch._dynamo.config
+    key = "recompile_limit" if hasattr(config, "recompile_limit") else "cache_size_limit"
+    old = getattr(config, key)
+    setattr(config, key, max(old, limit))
+    try:
+        yield
+    finally:
+        setattr(config, key, old)
+
+
+def _item_set(arr, idx, value, xp):
+    """
+    Assign ``value`` to ``arr[idx]`` and return the updated array.
+
+    JAX arrays are immutable and do not support in-place item assignment,
+    so the assignment is performed functionally using ``arr.at[idx].set``.
+    All other backends support in-place assignment directly.
+    """
+    if array_api_compat.is_jax_namespace(xp):
+        return arr.at[idx].set(value)
+    arr[idx] = value
+    return arr
+
+
+def _item_add(arr, idx, value, xp):
+    """
+    Add ``value`` to ``arr[idx]`` and return the updated array.
+
+    JAX arrays are immutable and do not support in-place item assignment,
+    so the addition is performed functionally using ``arr.at[idx].add``.
+    All other backends support in-place assignment directly.
+    """
+    if array_api_compat.is_jax_namespace(xp):
+        return arr.at[idx].add(value)
+    arr[idx] += value
+    return arr
+
+
 class BasisFunction:
     """
     Base class for all basis functions.
@@ -54,15 +108,15 @@ class BasisFunction:
     ...         # Change this if your new basis function is not in splinebox.basis_functions
     ...         return "splinebox.basis_functions.MyBasis()"
     ...
-    ...     def _func(self, t):
+    ...     def _func_numba(self, t):
     ...         # Implement your function here
     ...         return val
     ...
-    ...     def _derivative_1(self, t):
+    ...     def _derivative_1_numba(self, t):
     ...         # Implement the first derivative of your function here
     ...         return val
     ...
-    ...     def _derivative_2(self, t):
+    ...     def _derivative_2_numba(self, t):
     ...         # Implement the second derivative of your function here
     ...         # or raise an error if not differentiable
     ...         raise RuntimeError("MyBasis isn't twice differentiable.")
@@ -97,7 +151,7 @@ def __eq__(self, other):
             and other.support == self.support
         )
 
-    def __call__(self, t, derivative=0):
+    def __call__(self, t, derivative=0, jit=True):
         """
         Evaluate the function at position(s) t. You can optionally
         pass derivative=1 or 2 to obtain the respective derivative(s).
@@ -128,25 +182,51 @@ def __call__(self, t, derivative=0):
         >>> basis_function(t, derivative=1)
         array([ 1., -1.])
         """
+        xp = array_api_compat.numpy if isinstance(t, (int, float)) else array_api_compat.array_namespace(t)
+
+        use_numba = jit and array_api_compat.is_numpy_namespace(xp)
+
+        # Select the function that should be run
         if derivative == 0:
-            return self._func(t)
+            func = self._func_numba if use_numba else self._func_xp
         elif derivative == 1:
-            return self._derivative_1(t)
+            func = self._derivative_1_numba if use_numba else self._derivative_1_xp
         elif derivative == 2:
-            return self._derivative_2(t)
+            func = self._derivative_2_numba if use_numba else self._derivative_2_xp
         else:
             raise ValueError(f"derivative has to be 0, 1, or 2 not {derivative}")
 
-    def _func(self, t):
+        # Decide how to run the function
+        if use_numba:
+            return func(t)
+        elif jit and array_api_compat.is_torch_namespace(xp):
+            return xp.compile(fullgraph=True, dynamic=True, mode="reduce-overhead")(func)(t, xp)
+        elif jit and array_api_compat.is_jax_namespace(xp):
+            import jax
+
+            return jax.jit(func, static_argnames=["xp"])(t, xp)
+        else:
+            return func(t, xp)
+
+    def _func_numba(self, t):
+        raise NotImplementedError(BasisFunction._unimplemented_message)
+
+    def _func_xp(self, t, xp):
+        raise NotImplementedError(BasisFunction._unimplemented_message)
+
+    def _derivative_1_numba(self, t):
         raise NotImplementedError(BasisFunction._unimplemented_message)
 
-    def _derivative_1(self, t):
+    def _derivative_1_xp(self, t, xp):
         raise NotImplementedError(BasisFunction._unimplemented_message)
 
-    def _derivative_2(self, t):
+    def _derivative_2_numba(self, t):
         raise NotImplementedError(BasisFunction._unimplemented_message)
 
-    def filter_symmetric(self, s):
+    def _derivative_2_xp(self, t, xp):
+        raise NotImplementedError(BasisFunction._unimplemented_message)
+
+    def filter_symmetric(self, s, jit=True):
         """
         Returns a filtered version of the input s,
         used to convert knots into control points for an **open spline**.
@@ -191,9 +271,37 @@ def filter_symmetric(self, s):
 
         To perform the same test for the first and the last knot, padding is required.
         """
+        xp = array_api_compat.array_namespace(s)
+
+        s = xp.astype(s, float)
+
+        use_numba = jit and array_api_compat.is_numpy_namespace(xp)
+
+        if use_numba:
+            result = self._filter_symmetric_numba(s)
+        elif jit and array_api_compat.is_torch_namespace(xp):
+            with _torch_compile_recompile_limit():
+                result = xp.compile(fullgraph=True, dynamic=False, mode="reduce-overhead")(self._filter_symmetric_xp)(
+                    s, xp
+                )
+        elif jit and array_api_compat.is_jax_namespace(xp):
+            import jax
+
+            result = jax.jit(self._filter_symmetric_xp, static_argnames=["xp"])(s, xp)
+        else:
+            result = self._filter_symmetric_xp(s, xp)
+
+        if s.ndim == 1:
+            return xp.squeeze(result)
+        return result
+
+    def _filter_symmetric_numba(self, s):
+        raise NotImplementedError(BasisFunction._unimplemented_message)
+
+    def _filter_symmetric_xp(self, s, xp):
         raise NotImplementedError(BasisFunction._unimplemented_message)
 
-    def filter_periodic(self, s):
+    def filter_periodic(self, s, jit=True):
         """
         Returns a filtered version of the input s, used to convert
         knots into control points for a **closed spline**.
@@ -244,6 +352,34 @@ def filter_periodic(self, s):
         >>> b3(-1) * control_points[1] + b3(0) * control_points[2] + b3(1) * control_points[0]
         array([1., 3.])
         """
+        xp = array_api_compat.array_namespace(s)
+
+        s = xp.astype(s, float)
+
+        use_numba = jit and array_api_compat.is_numpy_namespace(xp)
+
+        if use_numba:
+            result = self._filter_periodic_numba(s)
+        elif jit and array_api_compat.is_torch_namespace(xp):
+            with _torch_compile_recompile_limit():
+                result = xp.compile(fullgraph=True, dynamic=False, mode="reduce-overhead")(self._filter_periodic_xp)(
+                    s, xp
+                )
+        elif jit and array_api_compat.is_jax_namespace(xp):
+            import jax
+
+            result = jax.jit(self._filter_periodic_xp, static_argnames=["xp"])(s, xp)
+        else:
+            result = self._filter_periodic_xp(s, xp)
+
+        if s.ndim == 1:
+            return xp.squeeze(result)
+        return result
+
+    def _filter_periodic_numba(self, s):
+        raise NotImplementedError(BasisFunction._unimplemented_message)
+
+    def _filter_periodic_xp(self, s, xp):
         raise NotImplementedError(BasisFunction._unimplemented_message)
 
     def refinement_mask(self):
@@ -312,15 +448,23 @@ def __repr__(self):
 
     @staticmethod
     @numba.vectorize([numba.float64(numba.float64)], nopython=True, cache=True)
-    def _func(t):  # pragma: no cover
+    def _func_numba(t):  # pragma: no cover
         val = 0
         if abs(t) >= 0 and abs(t) < 1:
             val = 1 - abs(t)
         return val
 
+    @staticmethod
+    def _func_xp(t, xp):  # pragma: no cover
+        return xp.where(
+            xp.abs(t) < 1,
+            1 - xp.abs(t),
+            0,
+        )
+
     @staticmethod
     @numba.vectorize([numba.float64(numba.float64)], nopython=True, cache=True)
-    def _derivative_1(t):  # pragma: no cover
+    def _derivative_1_numba(t):  # pragma: no cover
         val = 0
         if t > -1 and t < 0:
             val = 1
@@ -330,18 +474,37 @@ def _derivative_1(t):  # pragma: no cover
             val = np.nan
         return val
 
+    @staticmethod
+    def _derivative_1_xp(t, xp):  # pragma: no cover
+        abs_t = xp.abs(t)
+        result = -xp.sign(t)
+        result = xp.where((t == 0) | (abs_t == 1), xp.nan, result)
+        return xp.where(abs_t > 1, 0, result)
+
     @staticmethod
     @numba.vectorize([numba.float64(numba.float64)], nopython=True, cache=True)
-    def _derivative_2(t):
+    def _derivative_2_numba(t):
         return np.nan if t in (-1, 0, 1) else 0
 
     @staticmethod
-    def filter_symmetric(s):
-        return s.astype(float)
+    def _derivative_2_xp(t, xp):
+        return xp.where((t == -1) | (t == 0) | (t == 1), xp.nan, 0)
 
     @staticmethod
-    def filter_periodic(s):
-        return s.astype(float)
+    def _filter_symmetric_numba(s):
+        return s
+
+    @staticmethod
+    def _filter_symmetric_xp(s, xp):
+        return s
+
+    @staticmethod
+    def _filter_periodic_numba(s):
+        return s
+
+    @staticmethod
+    def _filter_periodic_xp(s, xp):
+        return s
 
     def refinement_mask(self):
         order = int(self.support)
@@ -403,7 +566,7 @@ def __repr__(self):
 
     @staticmethod
     @numba.vectorize([numba.float64(numba.float64)], nopython=True, cache=True)
-    def _func(t):  # pragma: no cover
+    def _func_numba(t):  # pragma: no cover
         val = 0
         if t >= -1.5 and t <= -0.5:
             val = 0.5 * (t**2) + 1.5 * t + 1.125
@@ -415,7 +578,7 @@ def _func(t):  # pragma: no cover
 
     @staticmethod
     @numba.vectorize([numba.float64(numba.float64)], nopython=True, cache=True)
-    def _derivative_1(t):  # pragma: no cover
+    def _derivative_1_numba(t):  # pragma: no cover
         val = 0
         if t >= -1.5 and t <= -0.5:
             val = t + 1.5
@@ -427,7 +590,7 @@ def _derivative_1(t):  # pragma: no cover
 
     @staticmethod
     @numba.vectorize([numba.float64(numba.float64)], nopython=True, cache=True)
-    def _derivative_2(t):  # pragma: no cover
+    def _derivative_2_numba(t):  # pragma: no cover
         val = np.nan
         if t < -1.5:
             val = 0
@@ -441,6 +604,62 @@ def _derivative_2(t):  # pragma: no cover
             val = 0
         return val
 
+    @staticmethod
+    def _func_xp(t, xp):  # pragma: no cover
+        return xp.where(
+            (t >= -1.5) & (t <= -0.5),
+            0.5 * (t**2) + 1.5 * t + 1.125,
+            xp.where(
+                (t > -0.5) & (t <= 0.5),
+                -t * t + 0.75,
+                xp.where(
+                    (t > 0.5) & (t <= 1.5),
+                    0.5 * (t**2) - 1.5 * t + 1.125,
+                    0,
+                ),
+            ),
+        )
+
+    @staticmethod
+    def _derivative_1_xp(t, xp):  # pragma: no cover
+        return xp.where(
+            (t >= -1.5) & (t <= -0.5),
+            t + 1.5,
+            xp.where(
+                (t > -0.5) & (t <= 0.5),
+                -2 * t,
+                xp.where(
+                    (t > 0.5) & (t <= 1.5),
+                    t - 1.5,
+                    0,
+                ),
+            ),
+        )
+
+    @staticmethod
+    def _derivative_2_xp(t, xp):  # pragma: no cover
+        return xp.where(
+            t < -1.5,
+            0,
+            xp.where(
+                (t > -1.5) & (t < -0.5),
+                1,
+                xp.where(
+                    (t > -0.5) & (t < 0.5),
+                    -2,
+                    xp.where(
+                        (t > 0.5) & (t < 1.5),
+                        1,
+                        xp.where(
+                            t > 1.5,
+                            0,
+                            xp.nan,
+                        ),
+                    ),
+                ),
+            ),
+        )
+
     def refinement_mask(self):
         order = int(self.support)
         mask = np.zeros(order + 1)
@@ -501,7 +720,7 @@ def __repr__(self):
 
     @staticmethod
     @numba.vectorize([numba.float64(numba.float64)], nopython=True, cache=True)
-    def _func(t):  # pragma: no cover
+    def _func_numba(t):  # pragma: no cover
         val = 0
         if abs(t) >= 0 and abs(t) < 1:
             val = 2 / 3 - (abs(t) ** 2) + (abs(t) ** 3) / 2
@@ -511,7 +730,7 @@ def _func(t):  # pragma: no cover
 
     @staticmethod
     @numba.vectorize([numba.float64(numba.float64)], nopython=True, cache=True)
-    def _derivative_1(t):  # pragma: no cover
+    def _derivative_1_numba(t):  # pragma: no cover
         val = 0
         if t >= 0 and t < 1:
             val = -2 * t + 1.5 * t * t
@@ -525,7 +744,7 @@ def _derivative_1(t):  # pragma: no cover
 
     @staticmethod
     @numba.vectorize([numba.float64(numba.float64)], nopython=True, cache=True)
-    def _derivative_2(t):  # pragma: no cover
+    def _derivative_2_numba(t):  # pragma: no cover
         val = 0
         if t >= 0 and t < 1:
             val = -2 + 3 * t
@@ -537,16 +756,62 @@ def _derivative_2(t):  # pragma: no cover
             val = 2 + t
         return val
 
-    def filter_symmetric(self, s):
-        # This is necessary for the docs to build correctly
-        result = self._filter_symmetric(s.astype(float))
-        if s.ndim == 1:
-            return np.squeeze(result)
-        return result
+    @staticmethod
+    def _func_xp(t, xp):  # pragma: no cover
+        abs_t = xp.abs(t)
+        return xp.where(
+            abs_t < 1,
+            2 / 3 - (abs_t**2) + (abs_t**3) / 2,
+            xp.where(
+                (abs_t >= 1) & (abs_t <= 2),
+                ((2 - abs_t) ** 3) / 6,
+                0,
+            ),
+        )
+
+    @staticmethod
+    def _derivative_1_xp(t, xp):  # pragma: no cover
+        return xp.where(
+            (t >= 0) & (t < 1),
+            -2 * t + 1.5 * t * t,
+            xp.where(
+                (t > -1) & (t < 0),
+                -2 * t - 1.5 * t * t,
+                xp.where(
+                    (t >= 1) & (t <= 2),
+                    -0.5 * ((2 - t) ** 2),
+                    xp.where(
+                        (t >= -2) & (t <= -1),
+                        0.5 * ((2 + t) ** 2),
+                        0,
+                    ),
+                ),
+            ),
+        )
+
+    @staticmethod
+    def _derivative_2_xp(t, xp):  # pragma: no cover
+        return xp.where(
+            (t >= 0) & (t < 1),
+            -2 + 3 * t,
+            xp.where(
+                (t > -1) & (t < 0),
+                -2 - 3 * t,
+                xp.where(
+                    (t >= 1) & (t <= 2),
+                    2 - t,
+                    xp.where(
+                        (t >= -2) & (t <= -1),
+                        2 + t,
+                        0,
+                    ),
+                ),
+            ),
+        )
 
     @staticmethod
     @numba.jit(nopython=True, nogil=True, cache=True)
-    def _filter_symmetric(s):  # pragma: no cover
+    def _filter_symmetric_numba(s):  # pragma: no cover
         M = len(s)
         pole = -2 + np.sqrt(3)
 
@@ -579,16 +844,39 @@ def _filter_symmetric(s):  # pragma: no cover
         c = c.reshape(shape)
         return c
 
-    def filter_periodic(self, s):
-        # This is necessary for the docs to build correctly
-        result = self._filter_periodic(s.astype(float))
-        if s.ndim == 1:
-            return np.squeeze(result)
-        return result
+    @staticmethod
+    def _filter_symmetric_xp(s, xp):  # pragma: no cover
+        M = len(s)
+        pole = -2 + math.sqrt(3)
+
+        ndim = 1 if s.ndim == 1 else s.shape[1]
+
+        cp = xp.zeros((M, ndim), dtype=s.dtype)
+        eps = 1e-8
+        k0 = min(((2 * M) - 2, int(math.ceil(math.log(eps) / math.log(abs(pole))))))
+        for k in range(k0):
+            m = k % (2 * M - 2)
+            val = s[2 * M - 2 - m] if m >= M else s[m]
+            cp = _item_add(cp, 0, val * (pole**m), xp)
+        cp = _item_set(cp, 0, cp[0] * (1 / (1 - (pole ** (2 * M - 2)))), xp)
+
+        for k in range(1, M):
+            cp = _item_set(cp, k, s[k] + pole * cp[k - 1], xp)
+
+        cm = xp.zeros((M, ndim), dtype=s.dtype)
+        cm = _item_set(cm, M - 1, cp[M - 1] + (pole * cp[M - 2]), xp)
+        cm = _item_set(cm, M - 1, cm[M - 1] * (pole / ((pole**2) - 1)), xp)
+
+        for k in range(M - 2, -1, -1):
+            cm = _item_set(cm, k, pole * (cm[k + 1] - cp[k]), xp)
+
+        c = cm * 6
+
+        return xp.where(xp.abs(c) < eps, 0, c)
 
     @staticmethod
     @numba.jit(nopython=True, nogil=True, cache=True)
-    def _filter_periodic(s):  # pragma: no cover
+    def _filter_periodic_numba(s):  # pragma: no cover
         M = len(s)
         pole = -2 + np.sqrt(3)
 
@@ -621,6 +909,36 @@ def _filter_periodic(s):  # pragma: no cover
         c = c.reshape(shape)
         return c
 
+    @staticmethod
+    def _filter_periodic_xp(s, xp):  # pragma: no cover
+        M = len(s)
+        pole = -2 + math.sqrt(3)
+
+        ndim = 1 if s.ndim == 1 else s.shape[1]
+
+        cp = xp.zeros((M, ndim), dtype=s.dtype)
+        for k in range(M):
+            cp = _item_add(cp, 0, s[(M - k) % M] * (pole**k), xp)
+        cp = _item_set(cp, 0, cp[0] * (1 / (1 - (pole**M))), xp)
+
+        for k in range(1, M):
+            cp = _item_set(cp, k, s[k] + pole * cp[k - 1], xp)
+
+        cm = xp.zeros((M, ndim), dtype=s.dtype)
+        for k in range(M):
+            cm = _item_add(cm, M - 1, (pole**k) * cp[k], xp)
+        cm = _item_set(cm, M - 1, cm[M - 1] * (pole / (1 - (pole**M))), xp)
+        cm = _item_add(cm, M - 1, cp[M - 1], xp)
+        cm = _item_set(cm, M - 1, cm[M - 1] * -pole, xp)
+
+        for k in range(M - 2, -1, -1):
+            cm = _item_set(cm, k, pole * (cm[k + 1] - cp[k]), xp)
+
+        c = cm * 6
+
+        eps = 1e-8
+        return xp.where(xp.abs(c) < eps, 0, c)
+
     def refinement_mask(self):
         order = int(self.support)
         mask = np.zeros(order + 1)
@@ -683,12 +1001,12 @@ def __repr__(self):
     def __eq__(self, other):
         return isinstance(other, type(self)) and other.M == self.M
 
-    def _func(self, t):
-        return self.__func(t, self.support / 2, self.M)
+    def _func_numba(self, t):
+        return self.__func_numba(t, self.support / 2, self.M)
 
     @staticmethod
     @numba.vectorize([numba.float64(numba.float64, numba.float64, numba.float64)], nopython=True, cache=True)
-    def __func(t, half_support, M):  # pragma: no cover
+    def __func_numba(t, half_support, M):  # pragma: no cover
         t += half_support
 
         alpha = np.pi / M
@@ -704,12 +1022,12 @@ def __func(t, half_support, M):  # pragma: no cover
 
         return L * val
 
-    def _derivative_1(self, t):
-        return self.__derivative_1(t, self.support / 2, self.M)
+    def _derivative_1_numba(self, t):
+        return self.__derivative_1_numba(t, self.support / 2, self.M)
 
     @staticmethod
     @numba.vectorize([numba.float64(numba.float64, numba.float64, numba.float64)], nopython=True, cache=True)
-    def __derivative_1(t, half_support, M):  # pragma: no cover
+    def __derivative_1_numba(t, half_support, M):  # pragma: no cover
         t += half_support
 
         alpha = np.pi / M
@@ -725,12 +1043,12 @@ def __derivative_1(t, half_support, M):  # pragma: no cover
 
         return L * val
 
-    def _derivative_2(self, t):
-        return self.__derivative_2(t, self.support / 2, self.M)
+    def _derivative_2_numba(self, t):
+        return self.__derivative_2_numba(t, self.support / 2, self.M)
 
     @staticmethod
     @numba.vectorize([numba.float64(numba.float64, numba.float64, numba.float64)], nopython=True, cache=True)
-    def __derivative_2(t, half_support, M):  # pragma: no cover
+    def __derivative_2_numba(t, half_support, M):  # pragma: no cover
         t += half_support
 
         alpha = np.pi / M
@@ -750,20 +1068,89 @@ def __derivative_2(t, half_support, M):  # pragma: no cover
 
         return L * val
 
-    def filter_symmetric(self, s):
-        self.M = len(s)
+    def _func_xp(self, t, xp):  # pragma: no cover
+        half_support = self.support / 2
+        t = t + half_support
+
+        alpha = xp.asarray(xp.pi / self.M)
+        L = 1 / (4 * xp.sin(alpha) ** 2)
+
+        val = xp.where(
+            (t >= 0) & (t < 1),
+            2 * xp.sin(alpha * t) ** 2,
+            xp.where(
+                (t >= 1) & (t < 2),
+                xp.cos(2 * alpha * (t - 2)) + xp.cos(2 * alpha * (t - 1)) - 2 * xp.cos(2 * alpha),
+                xp.where(
+                    (t >= 2) & (t <= 3),
+                    2 * xp.sin(alpha * (t - 3)) ** 2,
+                    0,
+                ),
+            ),
+        )
+        return L * val
 
-        b0 = self._func(0)
-        b1 = self._func(1)
+    def _derivative_1_xp(self, t, xp):  # pragma: no cover
+        half_support = self.support / 2
+        t = t + half_support
+
+        alpha = xp.asarray(xp.pi / self.M)
+        L = 1 / (4 * xp.sin(alpha) ** 2)
+
+        val = xp.where(
+            (t >= 0) & (t <= 1),
+            4 * alpha * xp.sin(alpha * t) * xp.cos(alpha * t),
+            xp.where(
+                (t > 1) & (t <= 2),
+                2 * alpha * (xp.sin(2 * alpha * (2 - t)) + xp.sin(2 * alpha * (1 - t))),
+                xp.where(
+                    (t > 2) & (t <= 3),
+                    4 * alpha * xp.sin(alpha * (t - 3)) * xp.cos(alpha * (t - 3)),
+                    0,
+                ),
+            ),
+        )
+        return L * val
 
-        result = self._filter_symmetric(s.astype(float), self.M, b0, b1)
-        if s.ndim == 1:
-            return np.squeeze(result)
-        return result
+    def _derivative_2_xp(self, t, xp):  # pragma: no cover
+        half_support = self.support / 2
+        t = t + half_support
+
+        alpha = xp.asarray(xp.pi / self.M)
+        L = 1 / (4 * xp.sin(alpha) ** 2)
+
+        val = xp.where(
+            t < 0,
+            0,
+            xp.where(
+                (t > 0) & (t < 1),
+                4 * alpha**2 * (xp.cos(alpha * t) ** 2 - xp.sin(alpha * t) ** 2),
+                xp.where(
+                    (t > 1) & (t < 2),
+                    -4 * alpha**2 * (xp.cos(2 * alpha * (2 - t)) + xp.cos(2 * alpha * (1 - t))),
+                    xp.where(
+                        (t > 2) & (t < 3),
+                        4 * alpha**2 * (xp.cos(alpha * (t - 3)) ** 2 - xp.sin(alpha * (t - 3)) ** 2),
+                        xp.where(
+                            t > 3,
+                            0,
+                            xp.nan,
+                        ),
+                    ),
+                ),
+            ),
+        )
+        return L * val
+
+    def _filter_symmetric_numba(self, s):
+        M = len(s)
+        b0 = self._func_numba(0)
+        b1 = self._func_numba(1)
+        return self.__filter_symmetric_numba(s, M, b0, b1)
 
     @staticmethod
     @numba.jit(nopython=True, nogil=True, cache=True)
-    def _filter_symmetric(s, M, b0, b1):  # pragma: no cover
+    def __filter_symmetric_numba(s, M, b0, b1):  # pragma: no cover
         ndim = 1 if s.ndim == 1 else s.shape[1]
         pole = (-b0 + np.sqrt(2 * b0 - 1)) / (1 - b0)
 
@@ -798,17 +1185,56 @@ def _filter_symmetric(s, M, b0, b1):  # pragma: no cover
         c = c.reshape(shape)
         return c
 
-    def filter_periodic(self, s):
-        self.M = len(s)
-        b0 = self._func(0)
-        result = self._filter_periodic(s.astype(float), self.M, b0)
-        if s.ndim == 1:
-            return np.squeeze(result)
-        return result
+    @staticmethod
+    def __b0_b1(M):
+        """
+        The values of the basis function at 0 and 1, computed with pure python
+        math so that the result can be compiled (e.g. with torch.compile).
+        """
+        alpha = math.pi / M
+        L = 1 / (4 * math.sin(alpha) ** 2)
+        b0 = L * (2 * math.cos(alpha) - 2 * math.cos(2 * alpha))
+        b1 = 2 * L * math.sin(alpha / 2) ** 2
+        return b0, b1
+
+    def _filter_symmetric_xp(self, s, xp):  # pragma: no cover
+        M = len(s)
+        b0, b1 = self.__b0_b1(self.M)
+        pole = (-b0 + math.sqrt(2 * b0 - 1)) / (1 - b0)
+
+        ndim = 1 if s.ndim == 1 else s.shape[1]
+
+        cp = xp.zeros((M, ndim), dtype=s.dtype)
+        eps = 1e-8
+        k0 = min(((2 * M) - 2, int(math.ceil(math.log(eps) / math.log(abs(pole))))))
+        for k in range(k0):
+            m = k % (2 * M - 2)
+            val = s[2 * M - 2 - m] if m >= M else s[m]
+            cp = _item_add(cp, 0, val * (pole**m), xp)
+        cp = _item_set(cp, 0, cp[0] * (1 / (1 - (pole ** (2 * M - 2)))), xp)
+
+        for k in range(1, M):
+            cp = _item_set(cp, k, s[k] + pole * cp[k - 1], xp)
+
+        cm = xp.zeros((M, ndim), dtype=s.dtype)
+        cm = _item_set(cm, M - 1, cp[M - 1] + (pole * cp[M - 2]), xp)
+        cm = _item_set(cm, M - 1, cm[M - 1] * (pole / ((pole * pole) - 1)), xp)
+
+        for k in range(M - 2, -1, -1):
+            cm = _item_set(cm, k, pole * (cm[k + 1] - cp[k]), xp)
+
+        c = cm / b1
+
+        return xp.where(xp.abs(c) < eps, 0, c)
+
+    def _filter_periodic_numba(self, s):
+        M = len(s)
+        b0 = self._func_numba(0)
+        return self.__filter_periodic_numba(s, M, b0)
 
     @staticmethod
     @numba.jit(nopython=True, nogil=True, cache=True)
-    def _filter_periodic(s, M, b0):  # pragma: no cover
+    def __filter_periodic_numba(s, M, b0):  # pragma: no cover
         ndim = 1 if s.ndim == 1 else s.shape[1]
 
         pole = (-b0 + np.sqrt(2 * b0 - 1)) / (1 - b0)
@@ -841,6 +1267,37 @@ def _filter_periodic(s, M, b0):  # pragma: no cover
         c = c.reshape(shape)
         return c
 
+    def _filter_periodic_xp(self, s, xp):  # pragma: no cover
+        M = len(s)
+        b0, _ = self.__b0_b1(self.M)
+        pole = (-b0 + math.sqrt(2 * b0 - 1)) / (1 - b0)
+
+        ndim = 1 if s.ndim == 1 else s.shape[1]
+
+        cp = xp.zeros((M, ndim), dtype=s.dtype)
+        cp = _item_set(cp, 0, s[0], xp)
+        for k in range(1, M):
+            cp = _item_add(cp, 0, s[k] * (pole ** (M - k)), xp)
+        cp = _item_set(cp, 0, cp[0] * (1 / (1 - (pole**M))), xp)
+
+        for k in range(1, M):
+            cp = _item_set(cp, k, s[k] + (pole * cp[k - 1]), xp)
+
+        cm = xp.zeros((M, ndim), dtype=s.dtype)
+        cm = _item_set(cm, M - 1, cp[M - 1], xp)
+        for k in range(M - 1):
+            cm = _item_add(cm, M - 1, cp[k] * (pole ** (k + 1)), xp)
+        cm = _item_set(cm, M - 1, cm[M - 1] * (1 / (1 - (pole**M))), xp)
+        cm = _item_set(cm, M - 1, cm[M - 1] * ((1 - pole) ** 2), xp)
+
+        for k in range(M - 2, -1, -1):
+            cm = _item_set(cm, k, (pole * cm[k + 1]) + (((1 - pole) ** 2) * cp[k]), xp)
+
+        c = cm
+
+        eps = 1e-8
+        return xp.where(xp.abs(c) < eps, 0, c)
+
     def refinement_mask(self):
         """
         TODO: This fails the tests. They are set to
@@ -915,7 +1372,7 @@ def __repr__(self):
 
     @staticmethod
     @numba.vectorize([numba.float64(numba.float64)], nopython=True, cache=True)
-    def _func(t):  # pragma: no cover
+    def _func_numba(t):  # pragma: no cover
         val = 0
         if np.abs(t) >= 0 and np.abs(t) <= 1:
             val = (3 / 2) * (np.abs(t) ** 3) - (5 / 2) * (np.abs(t) ** 2) + 1
@@ -925,7 +1382,7 @@ def _func(t):  # pragma: no cover
 
     @staticmethod
     @numba.vectorize([numba.float64(numba.float64)], nopython=True, cache=True)
-    def _derivative_1(t):  # pragma: no cover
+    def _derivative_1_numba(t):  # pragma: no cover
         val = 0
         if t >= 0 and t <= 1:
             val = t * (4.5 * t - 5)
@@ -939,7 +1396,7 @@ def _derivative_1(t):  # pragma: no cover
 
     @staticmethod
     @numba.vectorize([numba.float64(numba.float64)], nopython=True, cache=True)
-    def _derivative_2(t):  # pragma: no cover
+    def _derivative_2_numba(t):  # pragma: no cover
         val = 0
         if t >= 0 and t < 1:
             val = 9 * t - 5
@@ -954,12 +1411,81 @@ def _derivative_2(t):  # pragma: no cover
         return val
 
     @staticmethod
-    def filter_symmetric(s):
-        return s.astype(float)
+    def _func_xp(t, xp):  # pragma: no cover
+        abs_t = xp.abs(t)
+        return xp.where(
+            abs_t <= 1,
+            (3 / 2) * (abs_t**3) - (5 / 2) * (abs_t**2) + 1,
+            xp.where(
+                (abs_t > 1) & (abs_t <= 2),
+                (-1 / 2) * (abs_t**3) + (5 / 2) * (abs_t**2) - 4 * abs_t + 2,
+                0,
+            ),
+        )
 
     @staticmethod
-    def filter_periodic(s):
-        return s.astype(float)
+    def _derivative_1_xp(t, xp):  # pragma: no cover
+        return xp.where(
+            (t >= 0) & (t <= 1),
+            t * (4.5 * t - 5),
+            xp.where(
+                (t >= -1) & (t < 0),
+                -t * (4.5 * t + 5),
+                xp.where(
+                    (t > 1) & (t <= 2),
+                    -1.5 * t * t + 5 * t - 4,
+                    xp.where(
+                        (t >= -2) & (t < -1),
+                        1.5 * t * t + 5 * t + 4,
+                        0,
+                    ),
+                ),
+            ),
+        )
+
+    @staticmethod
+    def _derivative_2_xp(t, xp):  # pragma: no cover
+        return xp.where(
+            xp.abs(t) == 1,
+            xp.nan,
+            xp.where(
+                xp.abs(t) == 2,
+                xp.nan,
+                xp.where(
+                    (t >= 0) & (t < 1),
+                    9 * t - 5,
+                    xp.where(
+                        (t > -1) & (t < 0),
+                        -9 * t - 5,
+                        xp.where(
+                            (t > 1) & (t < 2),
+                            -3 * t + 5,
+                            xp.where(
+                                (t > -2) & (t < -1),
+                                3 * t + 5,
+                                0,
+                            ),
+                        ),
+                    ),
+                ),
+            ),
+        )
+
+    @staticmethod
+    def _filter_symmetric_numba(s):
+        return s
+
+    @staticmethod
+    def _filter_symmetric_xp(s, xp):
+        return s
+
+    @staticmethod
+    def _filter_periodic_numba(s):
+        return s
+
+    @staticmethod
+    def _filter_periodic_xp(s, xp):
+        return s
 
 
 class CubicHermite(BasisFunction):
@@ -1021,12 +1547,12 @@ def __str__(self):
     def __repr__(self):
         return "splinebox.basis_functions.CubicHermite()"
 
-    def _func(self, t):
-        return np.stack([self.h31(t), self.h32(t)], axis=-1)
+    def _func_numba(self, t):
+        return np.stack([self.h31_numba(t), self.h32_numba(t)], axis=-1)
 
     @staticmethod
     @numba.vectorize([numba.float64(numba.float64)], nopython=True, cache=True)
-    def h31(t):  # pragma: no cover
+    def h31_numba(t):  # pragma: no cover
         val = 0
         if t >= 0 and t <= 1:
             val = (1 + (2 * t)) * (t - 1) * (t - 1)
@@ -1036,7 +1562,7 @@ def h31(t):  # pragma: no cover
 
     @staticmethod
     @numba.vectorize([numba.float64(numba.float64)], nopython=True, cache=True)
-    def h32(t):  # pragma: no cover
+    def h32_numba(t):  # pragma: no cover
         val = 0
         if t >= 0 and t <= 1:
             val = t * (t - 1) * (t - 1)
@@ -1044,12 +1570,12 @@ def h32(t):  # pragma: no cover
             val = t * (t + 1) * (t + 1)
         return val
 
-    def _derivative_1(self, t):
-        return np.stack([self.h31prime(t), self.h32prime(t)], axis=-1)
+    def _derivative_1_numba(self, t):
+        return np.stack([self.h31prime_numba(t), self.h32prime_numba(t)], axis=-1)
 
     @staticmethod
     @numba.vectorize([numba.float64(numba.float64)], nopython=True, cache=True)
-    def h31prime(t):  # pragma: no cover
+    def h31prime_numba(t):  # pragma: no cover
         val = 0
         if t >= 0 and t <= 1:
             val = 6 * t * (t - 1)
@@ -1059,7 +1585,7 @@ def h31prime(t):  # pragma: no cover
 
     @staticmethod
     @numba.vectorize([numba.float64(numba.float64)], nopython=True, cache=True)
-    def h32prime(t):  # pragma: no cover
+    def h32prime_numba(t):  # pragma: no cover
         val = 0
         if t >= 0 and t <= 1:
             val = 3 * t * t - 4 * t + 1
@@ -1067,12 +1593,12 @@ def h32prime(t):  # pragma: no cover
             val = 3 * t * t + 4 * t + 1
         return val
 
-    def _derivative_2(self, t):
-        return np.stack([self.h31primeprime(t), self.h32primeprime(t)], axis=-1)
+    def _derivative_2_numba(self, t):
+        return np.stack([self.h31primeprime_numba(t), self.h32primeprime_numba(t)], axis=-1)
 
     @staticmethod
     @numba.vectorize([numba.float64(numba.float64)], nopython=True, cache=True)
-    def h31primeprime(t):  # pragma: no cover
+    def h31primeprime_numba(t):  # pragma: no cover
         val = 0
         if t >= 0 and t <= 1:
             val = 12 * t - 6
@@ -1082,7 +1608,7 @@ def h31primeprime(t):  # pragma: no cover
 
     @staticmethod
     @numba.vectorize([numba.float64(numba.float64)], nopython=True, cache=True)
-    def h32primeprime(t):  # pragma: no cover
+    def h32primeprime_numba(t):  # pragma: no cover
         val = 0
         if t >= 0 and t <= 1:
             val = 6 * t - 4
@@ -1090,6 +1616,87 @@ def h32primeprime(t):  # pragma: no cover
             val = 6 * t + 4
         return val
 
+    def _func_xp(self, t, xp):  # pragma: no cover
+        return xp.stack([self.h31_xp(t, xp), self.h32_xp(t, xp)], axis=-1)
+
+    @staticmethod
+    def h31_xp(t, xp):  # pragma: no cover
+        return xp.where(
+            (t >= 0) & (t <= 1),
+            (1 + (2 * t)) * (t - 1) * (t - 1),
+            xp.where(
+                (t < 0) & (t >= -1),
+                (1 - (2 * t)) * (t + 1) * (t + 1),
+                0,
+            ),
+        )
+
+    @staticmethod
+    def h32_xp(t, xp):  # pragma: no cover
+        return xp.where(
+            (t >= 0) & (t <= 1),
+            t * (t - 1) * (t - 1),
+            xp.where(
+                (t < 0) & (t >= -1),
+                t * (t + 1) * (t + 1),
+                0,
+            ),
+        )
+
+    def _derivative_1_xp(self, t, xp):  # pragma: no cover
+        return xp.stack([self.h31prime_xp(t, xp), self.h32prime_xp(t, xp)], axis=-1)
+
+    @staticmethod
+    def h31prime_xp(t, xp):  # pragma: no cover
+        return xp.where(
+            (t >= 0) & (t <= 1),
+            6 * t * (t - 1),
+            xp.where(
+                (t < 0) & (t >= -1),
+                -6 * t * (t + 1),
+                0,
+            ),
+        )
+
+    @staticmethod
+    def h32prime_xp(t, xp):  # pragma: no cover
+        return xp.where(
+            (t >= 0) & (t <= 1),
+            3 * t * t - 4 * t + 1,
+            xp.where(
+                (t < 0) & (t >= -1),
+                3 * t * t + 4 * t + 1,
+                0,
+            ),
+        )
+
+    def _derivative_2_xp(self, t, xp):  # pragma: no cover
+        return xp.stack([self.h31primeprime_xp(t, xp), self.h32primeprime_xp(t, xp)], axis=-1)
+
+    @staticmethod
+    def h31primeprime_xp(t, xp):  # pragma: no cover
+        return xp.where(
+            (t >= 0) & (t <= 1),
+            12 * t - 6,
+            xp.where(
+                (t < 0) & (t >= -1),
+                -12 * t - 6,
+                0,
+            ),
+        )
+
+    @staticmethod
+    def h32primeprime_xp(t, xp):  # pragma: no cover
+        return xp.where(
+            (t >= 0) & (t <= 1),
+            6 * t - 4,
+            xp.where(
+                (t < 0) & (t >= -1),
+                6 * t + 4,
+                0,
+            ),
+        )
+
     def h31_autocorrelation(self, i, j, M):  # pragma: no cover
         """
         Derived by V. Uhlmann during her PhD.
@@ -1202,12 +1809,20 @@ def h3_periodic_crosscorrelation(self, n, M):  # pragma: no cover
         return val
 
     @staticmethod
-    def filter_symmetric(s):
-        return s.astype(float)
+    def _filter_symmetric_numba(s):
+        return s
+
+    @staticmethod
+    def _filter_symmetric_xp(s, xp):
+        return s
 
     @staticmethod
-    def filter_periodic(s):
-        return s.astype(float)
+    def _filter_periodic_numba(s):
+        return s
+
+    @staticmethod
+    def _filter_periodic_xp(s, xp):
+        return s
 
 
 class ExponentialHermite(BasisFunction):
@@ -1273,12 +1888,12 @@ def __repr__(self):
     def __eq__(self, other):
         return isinstance(other, type(self)) and other.M == self.M
 
-    def _func(self, t):
-        return np.stack([self._he31(t, self.M), self._he32(t, self.M)], axis=-1)
+    def _func_numba(self, t):
+        return np.stack([self._he31_numba(t, self.M), self._he32_numba(t, self.M)], axis=-1)
 
     @staticmethod
     @numba.vectorize([numba.float64(numba.float64, numba.float64)], nopython=True, cache=True)
-    def _he31(t, M):  # pragma: no cover
+    def _he31_numba(t, M):  # pragma: no cover
         def _g1(t, M):
             val = 0
             if t >= 0 and t <= 1:
@@ -1297,7 +1912,7 @@ def _g1(t, M):
 
     @staticmethod
     @numba.vectorize([numba.float64(numba.float64, numba.float64)], nopython=True, cache=True)
-    def _he32(t, M):  # pragma: no cover
+    def _he32_numba(t, M):  # pragma: no cover
         def _g2(t, M):
             val = 0
             if t >= 0 and t <= 1:
@@ -1315,12 +1930,12 @@ def _g2(t, M):
         val = _g2(t, M) if t >= 0 else -1 * _g2(-t, M)
         return val
 
-    def _derivative_1(self, t):
-        return np.stack([self._he31prime(t, self.M), self._he32prime(t, self.M)], axis=-1)
+    def _derivative_1_numba(self, t):
+        return np.stack([self._he31prime_numba(t, self.M), self._he32prime_numba(t, self.M)], axis=-1)
 
     @staticmethod
     @numba.vectorize([numba.float64(numba.float64, numba.float64)], nopython=True, cache=True)
-    def _he31prime(t, M):  # pragma: no cover
+    def _he31prime_numba(t, M):  # pragma: no cover
         def _g1prime(t, M):
             val = 0
             if t >= 0 and t <= 1:
@@ -1335,7 +1950,7 @@ def _g1prime(t, M):
 
     @staticmethod
     @numba.vectorize([numba.float64(numba.float64, numba.float64)], nopython=True, cache=True)
-    def _he32prime(t, M):  # pragma: no cover
+    def _he32prime_numba(t, M):  # pragma: no cover
         def _g2prime(t, M):
             val = 0
             if t >= 0 and t <= 1:
@@ -1352,12 +1967,12 @@ def _g2prime(t, M):
         val = _g2prime(t, M) if t >= 0 else _g2prime(-t, M)
         return val
 
-    def _derivative_2(self, t):
-        return np.stack([self._he31primeprime(t, self.M), self._he32primeprime(t, self.M)], axis=-1)
+    def _derivative_2_numba(self, t):
+        return np.stack([self._he31primeprime_numba(t, self.M), self._he32primeprime_numba(t, self.M)], axis=-1)
 
     @staticmethod
     @numba.vectorize([numba.float64(numba.float64, numba.float64)], nopython=True, cache=True)
-    def _he31primeprime(t, M):  # pragma: no cover
+    def _he31primeprime_numba(t, M):  # pragma: no cover
         def _g1primeprime(t, M):
             val = 0
             if t >= 0 and t <= 1:
@@ -1372,7 +1987,7 @@ def _g1primeprime(t, M):
 
     @staticmethod
     @numba.vectorize([numba.float64(numba.float64, numba.float64)], nopython=True, cache=True)
-    def _he32primeprime(t, M):  # pragma: no cover
+    def _he32primeprime_numba(t, M):  # pragma: no cover
         def _g2primeprime(t, M):
             val = 0
             if t >= 0 and t <= 1:
@@ -1387,13 +2002,96 @@ def _g2primeprime(t, M):
         val = _g2primeprime(t, M) if t >= 0 else -1 * _g2primeprime(-t, M)
         return val
 
+    def _func_xp(self, t, xp):  # pragma: no cover
+        return xp.stack([self._he31_xp(t, self.M, xp), self._he32_xp(t, self.M, xp)], axis=-1)
+
+    @staticmethod
+    def _he31_xp(t, M, xp):  # pragma: no cover
+        abs_t = xp.abs(t)
+        alpha = xp.asarray(xp.pi / M)
+        denom = (alpha * xp.cos(alpha)) - xp.sin(alpha)
+        num = (
+            (0.5 * (2 * alpha * xp.cos(alpha) - xp.sin(alpha)))
+            - (alpha * xp.cos(alpha) * abs_t)
+            - (0.5 * xp.sin(alpha - (2 * alpha * abs_t)))
+        )
+        return xp.where((abs_t >= 0) & (abs_t <= 1), num / denom, 0)
+
+    @staticmethod
+    def _he32_xp(t, M, xp):  # pragma: no cover
+        abs_t = xp.abs(t)
+        alpha = xp.asarray(xp.pi / M)
+        denom = ((alpha * xp.cos(alpha)) - xp.sin(alpha)) * 8 * alpha * xp.sin(alpha)
+        num = (
+            -((2 * alpha * xp.cos(2 * alpha)) - xp.sin(2 * alpha))
+            - (4 * alpha * xp.sin(alpha) * xp.sin(alpha) * abs_t)
+            - (2 * xp.sin(alpha) * xp.cos(2 * alpha * (abs_t - 0.5)))
+            + (2 * alpha * xp.cos(2 * alpha * (abs_t - 1)))
+        )
+        val = xp.where((abs_t >= 0) & (abs_t <= 1), num / denom, 0)
+        return xp.where(t >= 0, val, -val)
+
+    def _derivative_1_xp(self, t, xp):  # pragma: no cover
+        return xp.stack([self._he31prime_xp(t, self.M, xp), self._he32prime_xp(t, self.M, xp)], axis=-1)
+
+    @staticmethod
+    def _he31prime_xp(t, M, xp):  # pragma: no cover
+        abs_t = xp.abs(t)
+        alpha = xp.asarray(xp.pi / M)
+        denom = (alpha * xp.cos(alpha)) - xp.sin(alpha)
+        num = -(alpha * xp.cos(alpha)) + (alpha * xp.cos(alpha - (2 * alpha * abs_t)))
+        val = xp.where((abs_t >= 0) & (abs_t <= 1), num / denom, 0)
+        return xp.where(t >= 0, val, -val)
+
+    @staticmethod
+    def _he32prime_xp(t, M, xp):  # pragma: no cover
+        abs_t = xp.abs(t)
+        alpha = xp.asarray(xp.pi / M)
+        denom = ((alpha * xp.cos(alpha)) - xp.sin(alpha)) * 8 * alpha * xp.sin(alpha)
+        num = (
+            -(4 * alpha * xp.sin(alpha) * xp.sin(alpha))
+            + (4 * alpha * xp.sin(alpha) * xp.sin(2 * alpha * (abs_t - 0.5)))
+            - (4 * alpha**2 * xp.sin(2 * alpha * (abs_t - 1)))
+        )
+        return xp.where((abs_t >= 0) & (abs_t <= 1), num / denom, 0)
+
+    def _derivative_2_xp(self, t, xp):  # pragma: no cover
+        return xp.stack([self._he31primeprime_xp(t, self.M, xp), self._he32primeprime_xp(t, self.M, xp)], axis=-1)
+
+    @staticmethod
+    def _he31primeprime_xp(t, M, xp):  # pragma: no cover
+        abs_t = xp.abs(t)
+        alpha = xp.asarray(xp.pi / M)
+        denom = (alpha * xp.cos(alpha)) - xp.sin(alpha)
+        num = 2 * alpha**2 * xp.sin(alpha - (2 * alpha * abs_t))
+        return xp.where((abs_t >= 0) & (abs_t <= 1), num / denom, 0)
+
+    @staticmethod
+    def _he32primeprime_xp(t, M, xp):  # pragma: no cover
+        abs_t = xp.abs(t)
+        alpha = xp.asarray(xp.pi / M)
+        denom = ((alpha * xp.cos(alpha)) - xp.sin(alpha)) * 8 * alpha * xp.sin(alpha)
+        num = +(8 * alpha**2 * xp.sin(alpha) * xp.cos(2 * alpha * (abs_t - 0.5))) - (
+            8 * alpha**3 * xp.cos(2 * alpha * (abs_t - 1))
+        )
+        val = xp.where((abs_t >= 0) & (abs_t <= 1), num / denom, 0)
+        return xp.where(t >= 0, val, -val)
+
+    @staticmethod
+    def _filter_symmetric_numba(s):
+        return s
+
+    @staticmethod
+    def _filter_symmetric_xp(s, xp):
+        return s
+
     @staticmethod
-    def filter_symmetric(s):
-        return s.astype(float)
+    def _filter_periodic_numba(s):
+        return s
 
     @staticmethod
-    def filter_periodic(s):
-        return s.astype(float)
+    def _filter_periodic_xp(s, xp):
+        return s
 
 
 def _multinomial(
diff --git a/src/splinebox/spline_curves.py b/src/splinebox/spline_curves.py
index 9c770dc..b1baef4 100644
--- a/src/splinebox/spline_curves.py
+++ b/src/splinebox/spline_curves.py
@@ -1339,7 +1339,7 @@ def curvature(self, t):
         >>> spline = splinebox.Spline(M=M, basis_function=splinebox.Exponential(M), closed=False)
         >>> spline.knots = np.array([[0, 0], [1, 1], [2, 0], [2.5, -0.5], [3, 0]])
         >>> spline.curvature([1, 3])
-        array([ 2.234, -4.104])
+        array([ 2.035, -3.616])
         """
         self._check_control_points()
         t, single_value = self._convert_to_array(t)
diff --git a/tests/conftest.py b/tests/conftest.py
index ab5a6d7..b25469d 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -1,13 +1,84 @@
+import contextlib
 import functools
 import itertools
 import math
 import unittest.mock
 
-import numpy as np
 import pytest
 import splinebox.basis_functions
 import splinebox.spline_curves
 
+ARRAY_BACKENDS_AND_DEVICES = []
+
+# NumPy
+with contextlib.suppress(ImportError):
+
+    import array_api_compat.numpy as np
+
+    ARRAY_BACKENDS_AND_DEVICES.append(pytest.param((np, "cpu"), id="numpy-cpu"))
+
+# PyTorch
+with contextlib.suppress(ImportError):
+
+    import array_api_compat.torch as torch
+    import torch as _torch
+
+    ARRAY_BACKENDS_AND_DEVICES.append(pytest.param((torch, "cpu"), id="torch-cpu"))
+
+    if _torch.cuda.is_available():
+        ARRAY_BACKENDS_AND_DEVICES.append(pytest.param((torch, "cuda"), id="torch-cuda"))
+
+# JAX
+with contextlib.suppress(ImportError):
+
+    import jax
+    import jax.numpy
+
+    # Enable 64-bit floats so that jax matches the float64 semantics of the
+    # other array backends (jax silently truncates float64 data otherwise).
+    jax.config.update("jax_enable_x64", True)
+
+    ARRAY_BACKENDS_AND_DEVICES.append(pytest.param((jax.numpy, jax.devices("cpu")[0]), id="jax-cpu"))
+
+    if any(device.platform == "gpu" for device in jax.devices()):
+        gpus = jax.devices("gpu")
+        ARRAY_BACKENDS_AND_DEVICES.append(pytest.param((jax.numpy, gpus[0]), id="jax-gpu"))
+
+
+@pytest.fixture(params=ARRAY_BACKENDS_AND_DEVICES)
+def xp_device(request):
+    return request.param
+
+
+@pytest.fixture
+def xp(xp_device):
+    return xp_device[0]
+
+
+@pytest.fixture
+def device(xp_device):
+    return xp_device[1]
+
+
+@pytest.fixture
+def as_xp_array(xp, device):
+    def _as_xp_array(obj, **kwargs):
+        if xp.__name__.endswith("numpy"):
+            return xp.asarray(obj, **kwargs)
+        return xp.asarray(obj, device=device, **kwargs)
+
+    return _as_xp_array
+
+
+@pytest.fixture(
+    params=[
+        # pytest.param(True, id="jit"),
+        pytest.param(False, id="no-jit"),
+    ]
+)
+def jit(request):
+    return request.param
+
 
 @pytest.fixture(
     params=[
@@ -97,12 +168,12 @@ def _is_locally_refinable(basis_function):
     return _is_locally_refinable
 
 
-@pytest.fixture(params=[4, 5, 8])
+@pytest.fixture(params=[4, 5, 8], ids=["M=4", "M=5", "M=8"])
 def M(request):
     return request.param
 
 
-@pytest.fixture(params=[True, False])
+@pytest.fixture(params=[True, False], ids=["closed", "open"])
 def closed(request):
     return request.param
 
@@ -173,7 +244,7 @@ def curvature_energy_compatible_spline_curve(curvature_energy_compatible_basis_f
     return splinebox.spline_curves.Spline(M, curvature_energy_compatible_basis_function, closed=False)
 
 
-@pytest.fixture(params=[0, 1, 2])
+@pytest.fixture(params=[0, 1, 2], ids=["derivative=0", "derivative=1", "derivative=2"])
 def derivative(request):
     return request.param
 
@@ -183,7 +254,7 @@ def call_positions(request):
     return request.param
 
 
-@pytest.fixture(params=[1, 2, 3, 5])
+@pytest.fixture(params=[1, 2, 3, 5], ids=["ndim=1", "ndim=2", "ndim=3", "ndim=5"])
 def codomain_dimensionality(request):
     """
     The number of dimensions the spline output
@@ -220,11 +291,11 @@ def control_point_gen_3D():
 
 
 @pytest.fixture
-def knot_gen(codomain_dimensionality):
+def knot_gen(codomain_dimensionality, as_xp_array):
     rng = np.random.default_rng(seed=2657)
 
     def _knot_gen(M=100):
-        return rng.random((M, codomain_dimensionality))
+        return as_xp_array(rng.random((M, codomain_dimensionality)))
 
     return _knot_gen
 
diff --git a/tests/test_basis_functions.py b/tests/test_basis_functions.py
index 5225308..e2b7451 100644
--- a/tests/test_basis_functions.py
+++ b/tests/test_basis_functions.py
@@ -5,25 +5,27 @@
 import splinebox.basis_functions
 
 
-def test_base_class_call(derivative):
+def test_base_class_call(derivative, xp, as_xp_array):
     basis_function = splinebox.basis_functions.BasisFunction(False, 2)
-    x = np.arange(10)
+    x = as_xp_array(xp.linspace(-4, 4, 10))
+
     with pytest.raises(NotImplementedError):
-        basis_function(x, derivative=derivative)
+        # The pytorch compiler does not work for raised errors.
+        basis_function(x, derivative=derivative, jit=False)
 
 
-def test_base_class_filters_and_refinement_mask():
+def test_base_class_filters_and_refinement_mask(xp, as_xp_array):
     basis_function = splinebox.basis_functions.BasisFunction(False, 2)
-    s = np.arange(10)
+    s = as_xp_array(xp.linspace(0, 4, 10))
     with pytest.raises(NotImplementedError):
-        basis_function.filter_symmetric(s)
+        basis_function.filter_symmetric(s, jit=False)
     with pytest.raises(NotImplementedError):
-        basis_function.filter_periodic(s)
+        basis_function.filter_periodic(s, jit=False)
     with pytest.raises(NotImplementedError):
         basis_function.refinement_mask()
 
 
-def test_filters(basis_function, is_interpolating, knot_gen, request):
+def test_filters(basis_function, is_interpolating, knot_gen, xp, request):
     if isinstance(basis_function, splinebox.basis_functions.B2):
         # The filter_symmetric and filter_periodic are not implemented for B2
         request.node.add_marker(pytest.mark.xfail)
@@ -32,12 +34,19 @@ def test_filters(basis_function, is_interpolating, knot_gen, request):
     if is_interpolating(basis_function):
         # The filters should not do anything since knot and control points are
         # the same for interpolating splines.
-        assert np.allclose(basis_function.filter_symmetric(s), s)
-        assert np.allclose(basis_function.filter_periodic(s), s)
+        assert xp.allclose(basis_function.filter_symmetric(s), s)
+        assert xp.allclose(basis_function.filter_periodic(s), s)
     else:
         # Flipping the knot or the control points should give the same result
-        assert np.allclose(basis_function.filter_symmetric(s), basis_function.filter_symmetric(s[::-1])[::-1])
-        assert np.allclose(basis_function.filter_periodic(s), basis_function.filter_periodic(s[::-1])[::-1])
+        flipped = xp.flip(s, axis=0)
+        assert xp.allclose(
+            basis_function.filter_symmetric(s),
+            xp.flip(basis_function.filter_symmetric(flipped), axis=0),
+        )
+        assert xp.allclose(
+            basis_function.filter_periodic(s),
+            xp.flip(basis_function.filter_periodic(flipped), axis=0),
+        )
 
 
 def test_refinement_mask(basis_function, is_locally_refinable, request):

From 5e0c079466cfbefe99c391fa33c8dd5d26c9bc4f Mon Sep 17 00:00:00 2001
From: faymanns <florian.aymanns@epfl.ch>
Date: Mon, 28 Sep 2026 15:29:04 +0200
Subject: [PATCH 02/15] Fixed wrong name for pytorch

---
 pyproject.toml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pyproject.toml b/pyproject.toml
index 1da2a06..c6fa985 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -95,7 +95,7 @@ test = [
     "pytest",
     "pytest-cov",
     "array-api-compat",
-    "pytorch",
+    "torch",
     "jax"
 ]
 examples = [

From f0c1cbb70e9b5c573e572259b3f1380f337d1c70 Mon Sep 17 00:00:00 2001
From: faymanns <florian.aymanns@epfl.ch>
Date: Mon, 28 Sep 2026 17:24:17 +0200
Subject: [PATCH 03/15] Stop pytest after first failed test

---
 tox.ini | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tox.ini b/tox.ini
index 7d2958f..d72aed8 100644
--- a/tox.ini
+++ b/tox.ini
@@ -27,5 +27,5 @@ passenv =
 extras =
     test,examples
 commands =
-    pytest -v --color=yes --doctest-modules src
-    pytest -v --color=yes --cov=splinebox --cov-report=xml
+    pytest -vx --color=yes --doctest-modules src
+    pytest -vx --color=yes --cov=splinebox --cov-report=xml

From 0e39e26bec308ef29af6c3db0bac2bcc1e39b7bf Mon Sep 17 00:00:00 2001
From: faymanns <florian.aymanns@epfl.ch>
Date: Mon, 28 Sep 2026 17:43:28 +0200
Subject: [PATCH 04/15] Pass environement variables for caching torch compiles

---
 .github/workflows/test_and_deploy.yml | 3 ++-
 tox.ini                               | 4 ++++
 2 files changed, 6 insertions(+), 1 deletion(-)

diff --git a/.github/workflows/test_and_deploy.yml b/.github/workflows/test_and_deploy.yml
index 8916816..9518967 100644
--- a/.github/workflows/test_and_deploy.yml
+++ b/.github/workflows/test_and_deploy.yml
@@ -45,9 +45,10 @@ jobs:
 
       # 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
diff --git a/tox.ini b/tox.ini
index d72aed8..39b60b9 100644
--- a/tox.ini
+++ b/tox.ini
@@ -24,6 +24,10 @@ platform =
 passenv =
     CI
     GITHUB_ACTIONS
+    USERNAME
+    USER
+    LOGNAME
+    TORCHINDUCTOR_CACHE_DIR
 extras =
     test,examples
 commands =

From c8ac18b3e4608e459a45fcc12c9ce83c1192d7e2 Mon Sep 17 00:00:00 2001
From: faymanns <florian.aymanns@epfl.ch>
Date: Mon, 28 Sep 2026 18:05:03 +0200
Subject: [PATCH 05/15] Add MSVC for windows runners

---
 .github/workflows/test_and_deploy.yml | 7 ++++++-
 tox.ini                               | 7 ++++++-
 2 files changed, 12 insertions(+), 2 deletions(-)

diff --git a/.github/workflows/test_and_deploy.yml b/.github/workflows/test_and_deploy.yml
index 9518967..b1fa2d4 100644
--- a/.github/workflows/test_and_deploy.yml
+++ b/.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,6 +43,11 @@ jobs:
           python -m pip install --upgrade pip
           python -m pip install setuptools tox tox-gh-actions
 
+      - name: Set up MSVC
+        uses: TheMrMilchmann/setup-msvc-dev@v4
+        with:
+          arch: x64
+
       # this runs the platform-specific tests declared in tox.ini
       - name: Test with tox
         env:
diff --git a/tox.ini b/tox.ini
index 39b60b9..50dd293 100644
--- a/tox.ini
+++ b/tox.ini
@@ -1,5 +1,5 @@
 [tox]
-envlist = py{39,310,311,312,313}-{linux,macos,windows}
+envlist = py{39,310,311,312,313,314}-{linux,macos,windows}
 isolated_build=true
 
 [gh-actions]
@@ -9,6 +9,7 @@ python =
     3.11: py311
     3.12: py312
     3.13: py313
+    3.14: py314
 
 [gh-actions:env]
 PLATFORM =
@@ -27,6 +28,10 @@ passenv =
     USERNAME
     USER
     LOGNAME
+    PATH
+    INCLUDE
+    LIB
+    LIBPATH
     TORCHINDUCTOR_CACHE_DIR
 extras =
     test,examples

From 7c331ed4207949fa4588a51560461b6fb093967b Mon Sep 17 00:00:00 2001
From: faymanns <florian.aymanns@epfl.ch>
Date: Mon, 28 Sep 2026 18:24:11 +0200
Subject: [PATCH 06/15] Only add MSVC on windows

---
 .github/workflows/test_and_deploy.yml | 1 +
 1 file changed, 1 insertion(+)

diff --git a/.github/workflows/test_and_deploy.yml b/.github/workflows/test_and_deploy.yml
index b1fa2d4..befa33e 100644
--- a/.github/workflows/test_and_deploy.yml
+++ b/.github/workflows/test_and_deploy.yml
@@ -44,6 +44,7 @@ jobs:
           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

From 0d5475cf1edc17069c5d800d1d4fdbb84dc66734 Mon Sep 17 00:00:00 2001
From: faymanns <florian.aymanns@epfl.ch>
Date: Mon, 28 Sep 2026 18:30:37 +0200
Subject: [PATCH 07/15] Cache compiled jax

---
 src/splinebox/basis_functions.py | 65 ++++++++++++++++++++++++++------
 1 file changed, 53 insertions(+), 12 deletions(-)

diff --git a/src/splinebox/basis_functions.py b/src/splinebox/basis_functions.py
index df1c981..d3dd400 100644
--- a/src/splinebox/basis_functions.py
+++ b/src/splinebox/basis_functions.py
@@ -45,6 +45,31 @@ def _torch_compile_recompile_limit(limit=64):
         setattr(config, key, old)
 
 
+def _cached_compiled(fn, key, compile_fn):
+    """
+    Return a compiled version of ``fn``, caching the compiled wrapper.
+
+    Constructing compiled wrappers, e.g. with ``torch.compile`` or
+    ``jax.jit``, is expensive, so the wrapper is constructed once per
+    function instead of on every call. The wrapper is cached on the
+    instance for bound methods, so that the ``self`` captured by the
+    cached wrapper matches the instance, and on the function itself
+    otherwise. ``key`` identifies the compilation configuration and must
+    be distinct for functions that are compiled differently, e.g. for
+    different backends.
+    """
+    if hasattr(fn, "__self__"):
+        owner, cache_key = fn.__self__, (fn.__func__, key)
+    elif hasattr(fn, "__func__"):
+        owner, cache_key = fn, (fn.__func__, key)
+    else:
+        owner, cache_key = fn, (fn, key)
+    cache = owner.__dict__.setdefault("_compiled_wrappers", {})
+    if cache_key not in cache:
+        cache[cache_key] = compile_fn(fn)
+    return cache[cache_key]
+
+
 def _item_set(arr, idx, value, xp):
     """
     Assign ``value`` to ``arr[idx]`` and return the updated array.
@@ -200,11 +225,13 @@ def __call__(self, t, derivative=0, jit=True):
         if use_numba:
             return func(t)
         elif jit and array_api_compat.is_torch_namespace(xp):
-            return xp.compile(fullgraph=True, dynamic=True, mode="reduce-overhead")(func)(t, xp)
+            return _cached_compiled(
+                func, "torch", lambda fn: xp.compile(fullgraph=True, dynamic=True, mode="reduce-overhead")(fn)
+            )(t, xp)
         elif jit and array_api_compat.is_jax_namespace(xp):
             import jax
 
-            return jax.jit(func, static_argnames=["xp"])(t, xp)
+            return _cached_compiled(func, "jax", lambda fn: jax.jit(fn, static_argnames=["xp"]))(t, xp)
         else:
             return func(t, xp)
 
@@ -281,18 +308,25 @@ def filter_symmetric(self, s, jit=True):
             result = self._filter_symmetric_numba(s)
         elif jit and array_api_compat.is_torch_namespace(xp):
             with _torch_compile_recompile_limit():
-                result = xp.compile(fullgraph=True, dynamic=False, mode="reduce-overhead")(self._filter_symmetric_xp)(
-                    s, xp
-                )
+                result = _cached_compiled(
+                    self._filter_symmetric_xp,
+                    "torch",
+                    lambda fn: xp.compile(fullgraph=True, dynamic=False, mode="reduce-overhead")(fn),
+                )(s, xp)
         elif jit and array_api_compat.is_jax_namespace(xp):
             import jax
 
-            result = jax.jit(self._filter_symmetric_xp, static_argnames=["xp"])(s, xp)
+            result = _cached_compiled(self._filter_symmetric_xp, "jax", lambda fn: jax.jit(fn, static_argnames=["xp"]))(
+                s, xp
+            )
         else:
             result = self._filter_symmetric_xp(s, xp)
 
         if s.ndim == 1:
-            return xp.squeeze(result)
+            # array_api_compat.torch requires the axis argument for squeeze,
+            # so the unit axes are determined explicitly.
+            axes = tuple(axis for axis, size in enumerate(result.shape) if size == 1)
+            return xp.squeeze(result, axis=axes) if axes else result
         return result
 
     def _filter_symmetric_numba(self, s):
@@ -362,18 +396,25 @@ def filter_periodic(self, s, jit=True):
             result = self._filter_periodic_numba(s)
         elif jit and array_api_compat.is_torch_namespace(xp):
             with _torch_compile_recompile_limit():
-                result = xp.compile(fullgraph=True, dynamic=False, mode="reduce-overhead")(self._filter_periodic_xp)(
-                    s, xp
-                )
+                result = _cached_compiled(
+                    self._filter_periodic_xp,
+                    "torch",
+                    lambda fn: xp.compile(fullgraph=True, dynamic=False, mode="reduce-overhead")(fn),
+                )(s, xp)
         elif jit and array_api_compat.is_jax_namespace(xp):
             import jax
 
-            result = jax.jit(self._filter_periodic_xp, static_argnames=["xp"])(s, xp)
+            result = _cached_compiled(self._filter_periodic_xp, "jax", lambda fn: jax.jit(fn, static_argnames=["xp"]))(
+                s, xp
+            )
         else:
             result = self._filter_periodic_xp(s, xp)
 
         if s.ndim == 1:
-            return xp.squeeze(result)
+            # array_api_compat.torch requires the axis argument for squeeze,
+            # so the unit axes are determined explicitly.
+            axes = tuple(axis for axis, size in enumerate(result.shape) if size == 1)
+            return xp.squeeze(result, axis=axes) if axes else result
         return result
 
     def _filter_periodic_numba(self, s):

From 9f08a0092e9922a5e917b03d4d1a072b59131b04 Mon Sep 17 00:00:00 2001
From: faymanns <florian.aymanns@epfl.ch>
Date: Tue, 29 Sep 2026 13:41:49 +0200
Subject: [PATCH 08/15] Module-level caching to avoid recompiles for non-static
 methods such as

---
 src/splinebox/basis_functions.py | 33 ++++++++++++++++++--------------
 1 file changed, 19 insertions(+), 14 deletions(-)

diff --git a/src/splinebox/basis_functions.py b/src/splinebox/basis_functions.py
index d3dd400..71c0b89 100644
--- a/src/splinebox/basis_functions.py
+++ b/src/splinebox/basis_functions.py
@@ -45,29 +45,34 @@ def _torch_compile_recompile_limit(limit=64):
         setattr(config, key, old)
 
 
+_compiled_wrapper_cache = {}
+
+
 def _cached_compiled(fn, key, compile_fn):
     """
     Return a compiled version of ``fn``, caching the compiled wrapper.
 
     Constructing compiled wrappers, e.g. with ``torch.compile`` or
-    ``jax.jit``, is expensive, so the wrapper is constructed once per
-    function instead of on every call. The wrapper is cached on the
-    instance for bound methods, so that the ``self`` captured by the
-    cached wrapper matches the instance, and on the function itself
-    otherwise. ``key`` identifies the compilation configuration and must
-    be distinct for functions that are compiled differently, e.g. for
+    ``jax.jit``, is expensive, so each wrapper is constructed only once.
+    Wrappers compiled for bound methods capture the instance they were
+    constructed for, so they can only be reused by instances whose
+    attributes are equal. The cache is therefore keyed by the instance's
+    attributes, which allows fresh instances with the same state, e.g. a
+    new ``Exponential(5)``, to reuse the wrapper compiled for a previous
+    one. This relies on basis function attributes not being modified after
+    the first compiled call, which holds because they are only set in
+    ``__init__``. Wrappers for plain functions are keyed by the function
+    itself. ``key`` identifies the compilation configuration and must be
+    distinct for functions that are compiled differently, e.g. for
     different backends.
     """
     if hasattr(fn, "__self__"):
-        owner, cache_key = fn.__self__, (fn.__func__, key)
-    elif hasattr(fn, "__func__"):
-        owner, cache_key = fn, (fn.__func__, key)
+        cache_key = (fn.__func__, key, tuple(sorted(fn.__self__.__dict__.items())))
     else:
-        owner, cache_key = fn, (fn, key)
-    cache = owner.__dict__.setdefault("_compiled_wrappers", {})
-    if cache_key not in cache:
-        cache[cache_key] = compile_fn(fn)
-    return cache[cache_key]
+        cache_key = (fn, key)
+    if cache_key not in _compiled_wrapper_cache:
+        _compiled_wrapper_cache[cache_key] = compile_fn(fn)
+    return _compiled_wrapper_cache[cache_key]
 
 
 def _item_set(arr, idx, value, xp):

From c6f4186e7a3e1955fbc3a5fc05896e96585eee25 Mon Sep 17 00:00:00 2001
From: faymanns <florian.aymanns@epfl.ch>
Date: Tue, 29 Sep 2026 17:00:21 +0200
Subject: [PATCH 09/15] Cache compiled methods on class-level

---
 src/splinebox/basis_functions.py | 44 +++++++++++++++++++-------------
 1 file changed, 26 insertions(+), 18 deletions(-)

diff --git a/src/splinebox/basis_functions.py b/src/splinebox/basis_functions.py
index 71c0b89..052b716 100644
--- a/src/splinebox/basis_functions.py
+++ b/src/splinebox/basis_functions.py
@@ -10,6 +10,7 @@
 For more information on implementing a new basis function, see :class:`splinebox.basis_functions.BasisFunction`.
 """
 
+import collections
 import contextlib
 import inspect
 import math
@@ -45,34 +46,40 @@ def _torch_compile_recompile_limit(limit=64):
         setattr(config, key, old)
 
 
-_compiled_wrapper_cache = {}
-
-
 def _cached_compiled(fn, key, compile_fn):
     """
     Return a compiled version of ``fn``, caching the compiled wrapper.
 
     Constructing compiled wrappers, e.g. with ``torch.compile`` or
     ``jax.jit``, is expensive, so each wrapper is constructed only once.
-    Wrappers compiled for bound methods capture the instance they were
-    constructed for, so they can only be reused by instances whose
-    attributes are equal. The cache is therefore keyed by the instance's
-    attributes, which allows fresh instances with the same state, e.g. a
-    new ``Exponential(5)``, to reuse the wrapper compiled for a previous
-    one. This relies on basis function attributes not being modified after
-    the first compiled call, which holds because they are only set in
-    ``__init__``. Wrappers for plain functions are keyed by the function
-    itself. ``key`` identifies the compilation configuration and must be
-    distinct for functions that are compiled differently, e.g. for
-    different backends.
+    The wrappers are cached in the ``_compiled_wrapper_cache`` dictionary
+    of the owning class. Wrappers compiled for bound methods capture the
+    instance they were compiled for, so the cache is keyed by the
+    instance's attributes in addition to the function. This allows fresh
+    instances with the same state, e.g. a new ``Exponential(5)``, to
+    reuse the wrapper compiled for a previous one, while instances with
+    different state, e.g. ``Exponential(7)``, get their own wrapper. It
+    relies on basis function attributes not being modified after the
+    first compiled call, which holds because they are only set in
+    ``__init__``. Note that this keeps the first instance with a given
+    state alive for the lifetime of the process. Wrappers for plain
+    functions are keyed by the function itself. ``key`` identifies the
+    compilation configuration and must be distinct for functions that are
+    compiled differently, e.g. for different backends.
     """
     if hasattr(fn, "__self__"):
-        cache_key = (fn.__func__, key, tuple(sorted(fn.__self__.__dict__.items())))
+        instance = fn.__self__
+        owner = type(instance)
+        cache_key = (fn.__func__, key, tuple(sorted(instance.__dict__.items())))
     else:
+        parts = fn.__qualname__.split(".")[:-1]  # drop the function name
+        owner = sys.modules[fn.__module__]
+        for part in parts:
+            owner = getattr(owner, part)
         cache_key = (fn, key)
-    if cache_key not in _compiled_wrapper_cache:
-        _compiled_wrapper_cache[cache_key] = compile_fn(fn)
-    return _compiled_wrapper_cache[cache_key]
+    if cache_key not in owner._compiled_wrapper_cache:
+        owner._compiled_wrapper_cache[cache_key] = compile_fn(fn)
+    return owner._compiled_wrapper_cache[cache_key]
 
 
 def _item_set(arr, idx, value, xp):
@@ -163,6 +170,7 @@ class BasisFunction:
     """
 
     _unimplemented_message = "This function is not implemented."
+    _compiled_wrapper_cache: dict[tuple, collections.abc.Callable] = {}
 
     def __init__(self, multigenerator, support):
         self.multigenerator = multigenerator

From 818ba5ea8f03a126e7cb114f7bf0843f29c76f3a Mon Sep 17 00:00:00 2001
From: faymanns <florian.aymanns@epfl.ch>
Date: Thu, 1 Oct 2026 11:51:47 +0200
Subject: [PATCH 10/15] spline __call__ working with pytorch

---
 src/splinebox/spline_curves.py | 137 +++++++++++++++++++--------------
 1 file changed, 81 insertions(+), 56 deletions(-)

diff --git a/src/splinebox/spline_curves.py b/src/splinebox/spline_curves.py
index b1baef4..7222688 100644
--- a/src/splinebox/spline_curves.py
+++ b/src/splinebox/spline_curves.py
@@ -9,6 +9,7 @@
 import math
 import warnings
 
+import array_api_compat
 import numba
 import numpy as np
 import scipy.integrate
@@ -154,7 +155,6 @@ class Spline:
     _wrong_array_size_msg = (
         "It looks like control_points is neither a 1 nor a 2D array. I don't know how to handle this yet."
     )
-    _no_control_points_msg = "This spline object doesn't have any control points yet."
     _unimplemented_msg = "This function is not implemented."
 
     def __init__(
@@ -188,14 +188,16 @@ def __init__(
         self._cached_segments = None
         self._cached_segment_lengths = None
 
-    def _check_control_points(self):
+    def _get_control_points_xp(self):
         """
         Most methods require control points to be set before they
         can be used. This helper function checks if the control points have been
-        set.
+        set and return the their array namespace.
         """
         if self.control_points is None:
-            raise RuntimeError(self._no_control_points_msg)
+            raise RuntimeError("This spline object doesn't have any control points yet.")
+        else:
+            return array_api_compat.array_namespace(self.control_points)
 
     def __str__(self):
         closed_str = "closed" if self.closed else "open"
@@ -231,7 +233,8 @@ def control_points(self):
     @control_points.setter
     def control_points(self, values):
         if values is not None:
-            values = np.array(values)
+            if not array_api_compat.is_array_api_obj(values):
+                values = np.array(values)
             if values.ndim != 2:
                 raise ValueError(
                     "The control point array has to be 2D. The first dimension encodes the control points and the second corresponds to the dimensionality of the codomain of the spline."
@@ -796,7 +799,7 @@ def draw(self, x, y):
         >>> plt.plot(vals[:, 0], vals[:, 1])  # doctest: +SKIP
         >>> plt.show()  # doctest: +SKIP
         """
-        self._check_control_points()
+        self._get_control_points_xp()
 
         if self.ndim != 2:
             raise RuntimeError("draw() can only be used with 2D curves")
@@ -822,7 +825,7 @@ def dtheta(self, t):
         .. math::
             \frac{d \theta}{dt} = \frac{1}{r^2} \left( x\frac{dy}{dt} - y\frac{dx}{dt} \right) \text{, where } r^2 = x^2 + y^2
         """
-        self._check_control_points()
+        self._get_control_points_xp()
         if self.ndim != 2:
             raise RuntimeError("dtheta() is only defined for 2D curves.")
         t, single_value = self._convert_to_array(t)
@@ -891,7 +894,7 @@ def is_inside(self, x, y):
         """
         if not self.closed:
             raise RuntimeError("isInside() can only be used with closed curves.")
-        self._check_control_points()
+        self._get_control_points_xp()
         if self.ndim != 2:
             raise RuntimeError("isInside() can only be used with 2D curves.")
 
@@ -1090,7 +1093,7 @@ def arc_length(self, stop=None, start=0):
         >>> spline.arc_length(np.arange(M - 1), np.arange(1, M))
         array([4.982, 5.288, 3.09 , 3.675])
         """
-        self._check_control_points()
+        self._get_control_points_xp()
 
         if stop is None:
             stop = self.M if self.closed else self.M - 1
@@ -1209,7 +1212,7 @@ def arc_length_to_parameter(self, s, atol=1e-4):
         >>> spline.arc_length(0, 2.12)  # doctest: +NUMBER
         2.2
         """
-        self._check_control_points()
+        self._get_control_points_xp()
         s, single_value = self._convert_to_array(s)
         sort_indices = np.argsort(s)
         results = np.zeros_like(s, dtype=float)
@@ -1341,7 +1344,7 @@ def curvature(self, t):
         >>> spline.curvature([1, 3])
         array([ 2.035, -3.616])
         """
-        self._check_control_points()
+        self._get_control_points_xp()
         t, single_value = self._convert_to_array(t)
         first_deriv = self(t, derivative=1)
         second_deriv = self(t, derivative=2)
@@ -1414,7 +1417,7 @@ def normal(self, t, frame="bishop", initial_vector=None, shift=None):
         >>> plt.arrow(point[0], point[1], normal[0], normal[1])  # doctest: +SKIP
         >>> plt.show()  # doctest: +SKIP
         """
-        self._check_control_points()
+        self._get_control_points_xp()
         t, single_value = self._convert_to_array(t)
         if self.ndim not in (2, 3):
             raise NotImplementedError(
@@ -1537,7 +1540,7 @@ def moving_frame(self, t, method="frenet", initial_vector=None, shift=None):
                 [-1.,  0.,  0.],
                 [ 0., -0.,  1.]]])
         """
-        self._check_control_points()
+        self._get_control_points_xp()
 
         if self.ndim != 3:
             raise RuntimeError("A frame can only be computed for splines in 3D.")
@@ -1656,7 +1659,7 @@ def moving_frame(self, t, method="frenet", initial_vector=None, shift=None):
         nopython=True,
         cache=True,
     )
-    def _compute_tval_and_indices(t, shift, closed, M, pad, tval, indices):
+    def _compute_tval_and_indices_numba(t, shift, closed, M, pad, tval, indices):
         """
         Helper function for creating a csr sparse matrix.
 
@@ -1686,6 +1689,14 @@ def _compute_tval_and_indices(t, shift, closed, M, pad, tval, indices):
         else:
             indices[:] = (t - t_mod_1)[:, np.newaxis] + shift[np.newaxis, :] + pad
 
+    @staticmethod
+    def _compute_tval_and_indices_xp(t, shift, closed, M, pad, xp):
+        t_mod_1 = t % 1
+        tval = t_mod_1[:, xp.newaxis] - shift[xp.newaxis, :]
+        indices = xp.asarray(t - t_mod_1, dtype=int)[:, xp.newaxis] + shift[xp.newaxis, :]
+        indices = indices % M if closed else indices + pad
+        return tval, indices
+
     def basis_matrix(self, t, derivative=0):
         r"""
         Computes the basis matrix :math:`\mathbf{\Phi}` as defined in :ref:`theory/data_approximation:Data approximation`.
@@ -1703,10 +1714,11 @@ def basis_matrix(self, t, derivative=0):
         bm : scipy.sparse.csr_array
             The basis matrix.
         """
-        t, single_value = self._convert_to_array(t)
+        xp = self._get_control_points_xp()
+        t, single_value = self._convert_to_array(t, xp)
 
-        minimum = np.min(t)
-        maximum = np.max(t)
+        minimum = xp.min(t)
+        maximum = xp.max(t)
         if np.isnan(maximum):
             raise ValueError("t should not contain any NaN values.")
 
@@ -1714,38 +1726,45 @@ def basis_matrix(self, t, derivative=0):
         shift = np.arange(-bound + 1, bound + 1)
         n_control_points = self.M if self.closed else self.M + 2 * self.pad
 
-        tval = np.empty((len(t), len(shift)), dtype=float)
-        indices = np.empty((len(t), len(shift)), dtype=int)
-        self._compute_tval_and_indices(t, shift, self.closed, self.M, self.pad, tval, indices)
-        data = self.basis_function(tval, derivative=derivative)
+        if array_api_compat.is_numpy_namespace(xp):
+            tval = xp.empty((len(t), len(shift)), dtype=float)
+            indices = xp.empty((len(t), len(shift)), dtype=int)
+            self._compute_tval_and_indices_numba(t, shift, self.closed, self.M, self.pad, tval, indices)
 
-        if not self.closed and (minimum < 0 or maximum > self.M - 1):
-            mask = (indices >= 0) & (indices <= (self.M - 1 + 2 * self.pad))
+            data = self.basis_function(tval, derivative=derivative)
 
-            # Mask rows that are completely empty
-            row_mask = np.any(mask, axis=1)
-            data = data[row_mask]
-            indices = indices[row_mask]
+            if not self.closed and (minimum < 0 or maximum > self.M - 1):
+                mask = (indices >= 0) & (indices <= (self.M - 1 + 2 * self.pad))
 
-            mask = mask[row_mask]
-            data[~mask] = 0
+                # Mask rows that are completely empty
+                row_mask = np.any(mask, axis=1)
+                data = data[row_mask]
+                indices = indices[row_mask]
 
-            data = data.flatten()
-            indices = indices.flatten()
+                mask = mask[row_mask]
+                data[~mask] = 0
 
-            indptr = np.full(len(t) + 1, len(indices))
-            indptr[: len(indices) // len(shift) + 1] = np.arange(0, len(indices) + 1, len(shift))
+                data = data.flatten()
+                indices = indices.flatten()
+
+                indptr = np.full(len(t) + 1, len(indices))
+                indptr[: len(indices) // len(shift) + 1] = np.arange(0, len(indices) + 1, len(shift))
 
+            else:
+                data = data.flatten()
+                indices = indices.flatten()
+                indptr = np.arange(0, len(indices) + 1, len(shift))
+
+            # This is necessary because there is not enough padding of the control points.
+            # When t=M-1 shift reaches past the last padded control point.
+            # The seems to be the limiting factor. An alternative solution should be found.
+            bm = scipy.sparse.csr_array((data, indices, indptr), shape=(len(t), n_control_points + 1), copy=False)
+            bm = bm[:, :n_control_points]
         else:
-            data = data.flatten()
-            indices = indices.flatten()
-            indptr = np.arange(0, len(indices) + 1, len(shift))
+            tval, indices = self._compute_tval_and_indices_xp(t, shift, self.closed, self.M, self.pad, xp)
+
+            basis_function_values = self.basis_function(tval, derivative=derivative)
 
-        # This is necessary because there is not enough padding of the control points.
-        # When t=M-1 shift reaches past the last padded control point.
-        # The seems to be the limiting factor. An alternative solution should be found.
-        bm = scipy.sparse.csr_array((data, indices, indptr), shape=(len(t), n_control_points + 1), copy=False)
-        bm = bm[:, :n_control_points]
         return bm
 
     def __call__(self, t, derivative=0):
@@ -1780,8 +1799,8 @@ def __call__(self, t, derivative=0):
                [ 1.5,  0.5],
                [ 3. ,  1. ]])
         """
-        self._check_control_points()
-        t, single_value = self._convert_to_array(t)
+        xp = self._get_control_points_xp()
+        t, single_value = self._convert_to_array(t, xp)
 
         bm = self.basis_matrix(t, derivative=derivative)
 
@@ -1891,7 +1910,7 @@ def derivative_of_curvature_energy_wrt_control_points(self):
         """
         return 2 * np.einsum("mx,lm->lx", self.control_points, self._h3)
 
-    def _convert_to_array(self, t):
+    def _convert_to_array(self, t, xp):
         """
         Helper function that converts a function input
         to an array. This allows functions to accept int and float
@@ -1902,13 +1921,19 @@ def _convert_to_array(self, t):
         """
         single_value = False
         if not isinstance(t, collections.abc.Iterable):
-            t = np.array([t])
+            t = xp.asarray([t])
             single_value = True
-        elif isinstance(t, np.ndarray) and t.shape == ():
+        elif not array_api_compat.is_array_api_obj(t):
+            t = xp.asarray(t)
+        elif array_api_compat.array_namespace(t) is not xp:
+            raise ValueError(
+                f"The control points are a {xp.__name__.split('.')[-1]} array but "
+                "the parameter values you provided are a "
+                f"{array_api_compat.array_namespace(t).__name__.split('.')[-1]} array."
+            )
+        elif t.shape == ():
             # Array with only one element, e.g. np.array(0.)
-            t = np.array([t.item()])
-        elif not isinstance(t, np.ndarray):
-            t = np.array(t)
+            t = xp.asarray([t.item()])
         if t.ndim > 1:
             raise ValueError("The parameter array has to be 1D.")
         return t, single_value
@@ -1919,7 +1944,7 @@ def _control_points_centroid(self):
         and :meth:`splinebox.spline_curves.Spline.rotate`.
         Computes the centroid of the coefficients.
         """
-        self._check_control_points()
+        self._get_control_points_xp()
         return np.mean(self.control_points, axis=0)
 
     def translate(self, vector):
@@ -1957,7 +1982,7 @@ def translate(self, vector):
         >>> plt.legend()  # doctest: +SKIP
         >>> plt.show()  # doctest: +SKIP
         """
-        self._check_control_points()
+        self._get_control_points_xp()
         self.control_points = self.control_points + vector
 
     def scale(self, scaling_factor):
@@ -1995,7 +2020,7 @@ def scale(self, scaling_factor):
         >>> plt.legend()  # doctest: +SKIP
         >>> plt.show()  # doctest: +SKIP
         """
-        self._check_control_points()
+        self._get_control_points_xp()
         centroid = self._control_points_centroid()
         self.translate(-centroid)
         self.control_points *= scaling_factor
@@ -2038,7 +2063,7 @@ def rotate(self, rotation_matrix, centred=True):
         >>> plt.legend()  # doctest: +SKIP
         >>> plt.show()  # doctest: +SKIP
         """
-        self._check_control_points()
+        self._get_control_points_xp()
         if self.ndim == 1:
             raise RuntimeError("1D splines can not be rotated.")
 
@@ -2104,7 +2129,7 @@ def distance(self, points, return_t=False):
         >>> plt.gca().set_aspect("equal")  # doctest: +SKIP
         >>> plt.show()  # doctest: +SKIP
         """
-        self._check_control_points()
+        self._get_control_points_xp()
         if self.ndim == 1:
             raise RuntimeError("Cannot compute distance for 1D splines.")
 
@@ -2302,7 +2327,7 @@ def mesh(
         >>> connectivity.shape
         (6480, 4)
         """
-        self._check_control_points()
+        self._get_control_points_xp()
         if self.ndim != 3:
             raise NotImplementedError("Meshes are only implemented for splines in 3D.")
         cap_ends = self._normalize_cap_ends(cap_ends) if mesh_type == "surface" and not self.closed else None
@@ -2575,7 +2600,7 @@ def _check_control_points_and_tangents(self):
         can be used. This helper function checks if control points and tangents have been
         set.
         """
-        self._check_control_points()
+        self._get_control_points_xp()
         if self.tangents is None:
             raise RuntimeError(self._no_tangents_msg)
 

From c1d89d299a52361258819890714649bee1156f01 Mon Sep 17 00:00:00 2001
From: faymanns <florian.aymanns@epfl.ch>
Date: Sun, 11 Oct 2026 20:22:10 +0200
Subject: [PATCH 11/15] Test for different backends with fixtures

---
 tests/conftest.py | 110 +++++++++++++++++++++++++++-------------------
 1 file changed, 65 insertions(+), 45 deletions(-)

diff --git a/tests/conftest.py b/tests/conftest.py
index b25469d..364eb16 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -1,3 +1,4 @@
+import collections
 import contextlib
 import functools
 import itertools
@@ -23,10 +24,10 @@
     import array_api_compat.torch as torch
     import torch as _torch
 
-    ARRAY_BACKENDS_AND_DEVICES.append(pytest.param((torch, "cpu"), id="torch-cpu"))
+    ARRAY_BACKENDS_AND_DEVICES.append(pytest.param((torch, torch.device("cpu")), id="torch-cpu"))
 
     if _torch.cuda.is_available():
-        ARRAY_BACKENDS_AND_DEVICES.append(pytest.param((torch, "cuda"), id="torch-cuda"))
+        ARRAY_BACKENDS_AND_DEVICES.append(pytest.param((torch, torch.device("cuda")), id="torch-cuda"))
 
 # JAX
 with contextlib.suppress(ImportError):
@@ -63,8 +64,6 @@ def device(xp_device):
 @pytest.fixture
 def as_xp_array(xp, device):
     def _as_xp_array(obj, **kwargs):
-        if xp.__name__.endswith("numpy"):
-            return xp.asarray(obj, **kwargs)
         return xp.asarray(obj, device=device, **kwargs)
 
     return _as_xp_array
@@ -184,18 +183,18 @@ def padding_function(request):
 
 
 @pytest.fixture
-def spline_curve(basis_function, M, closed, padding_function):
+def spline_curve(basis_function, M, closed, padding_function, xp, device):
     if isinstance(
         basis_function, (splinebox.basis_functions.CubicHermite, splinebox.basis_functions.ExponentialHermite)
     ):
         spline = splinebox.spline_curves.HermiteSpline(
             M, basis_function, closed=closed, padding_function=padding_function
         )
-        spline._check_control_points_and_tangents = unittest.mock.MagicMock()
-        spline._check_control_points = unittest.mock.MagicMock()
+        spline._get_control_points_and_tangents_xp_and_device = unittest.mock.MagicMock(return_value=(xp, device))
+        spline._get_control_points_xp_and_device = unittest.mock.MagicMock(return_value=(xp, device))
     else:
         spline = splinebox.spline_curves.Spline(M, basis_function, closed=closed, padding_function=padding_function)
-        spline._check_control_points = unittest.mock.MagicMock()
+        spline._get_control_points_xp_and_device = unittest.mock.MagicMock(return_value=(xp, device))
     return spline
 
 
@@ -250,7 +249,25 @@ def derivative(request):
 
 
 @pytest.fixture(params=[1.5, np.array(0.0), np.linspace(0, 10, 1000)])
-def call_positions(request):
+def call_positions(request, as_xp_array):
+    if isinstance(request.param, collections.abc.Iterable):
+        return as_xp_array(request.param)
+    return request.param
+
+
+@pytest.fixture(
+    params=[
+        0.5,
+        1,
+        [0.5],
+        [0, 1],
+        [-0.5, 0.0, 0.5],
+        [[-0.7, -0.2, 0.3], [-0.6, -0.1, 0.4], [-0.5, 0.0, 0.5], [-0.4, 0.1, 0.6]],
+    ]
+)
+def basis_function_call_positions(request, as_xp_array):
+    if isinstance(request.param, collections.abc.Iterable):
+        return as_xp_array(request.param)
     return request.param
 
 
@@ -265,29 +282,29 @@ def codomain_dimensionality(request):
     return request.param
 
 
-def _control_point_gen(M, support, closed, ndim, rng=None):
+def _control_point_gen(M, support, closed, ndim, as_xp_array, rng=None):
     if rng is None:
         rng = np.random.default_rng(seed=42)
     points = rng.random((M, ndim)) * 100 if closed else rng.random((M + 2 * (math.ceil(support / 2) - 1), ndim)) * 100
-    return points
+    return as_xp_array(points)
 
 
 @pytest.fixture
-def coeff_gen(codomain_dimensionality):
+def coeff_gen(codomain_dimensionality, as_xp_array):
     rng = np.random.default_rng(seed=1492)
-    return functools.partial(_control_point_gen, ndim=codomain_dimensionality, rng=rng)
+    return functools.partial(_control_point_gen, ndim=codomain_dimensionality, as_xp_array=as_xp_array, rng=rng)
 
 
 @pytest.fixture
-def control_point_gen_2D():
+def control_point_gen_2D(as_xp_array):
     rng = np.random.default_rng(seed=9052)
-    return functools.partial(_control_point_gen, ndim=2, rng=rng)
+    return functools.partial(_control_point_gen, ndim=2, as_xp_array=as_xp_array, rng=rng)
 
 
 @pytest.fixture
-def control_point_gen_3D():
+def control_point_gen_3D(as_xp_array):
     rng = np.random.default_rng(seed=7755)
-    return functools.partial(_control_point_gen, ndim=3, rng=rng)
+    return functools.partial(_control_point_gen, ndim=3, as_xp_array=as_xp_array, rng=rng)
 
 
 @pytest.fixture
@@ -334,7 +351,7 @@ def translation_vector(codomain_dimensionality):
     return np.random.rand(codomain_dimensionality) * 10 - 5
 
 
-def _initialize_spline_curve(spline, control_point_gen):
+def _initialise_spline_curve(spline, control_point_gen):
     support = spline.basis_function.support
     closed = spline.closed
     M = spline.M
@@ -347,58 +364,58 @@ def _initialize_spline_curve(spline, control_point_gen):
 
 
 @pytest.fixture
-def initialize_spline_curve(is_hermite_spline, coeff_gen):
-    return functools.partial(_initialize_spline_curve, control_point_gen=coeff_gen)
+def initialise_spline_curve(is_hermite_spline, coeff_gen):
+    return functools.partial(_initialise_spline_curve, control_point_gen=coeff_gen)
 
 
 @pytest.fixture
-def initialize_2D_spline_curve(is_hermite_spline, control_point_gen_2D):
-    return functools.partial(_initialize_spline_curve, control_point_gen=control_point_gen_2D)
+def initialise_2D_spline_curve(is_hermite_spline, control_point_gen_2D):
+    return functools.partial(_initialise_spline_curve, control_point_gen=control_point_gen_2D)
 
 
 @pytest.fixture
-def initialize_3D_spline_curve(is_hermite_spline, control_point_gen_3D):
-    return functools.partial(_initialize_spline_curve, control_point_gen=control_point_gen_3D)
+def initialise_3D_spline_curve(is_hermite_spline, control_point_gen_3D):
+    return functools.partial(_initialise_spline_curve, control_point_gen=control_point_gen_3D)
 
 
 @pytest.fixture
-def initialized_spline_curve(spline_curve, initialize_spline_curve):
-    return initialize_spline_curve(spline_curve)
+def initialised_spline_curve(spline_curve, initialise_spline_curve):
+    return initialise_spline_curve(spline_curve)
 
 
 @pytest.fixture
-def initialized_twice_differentiable_spline_curve(twice_differentiable_spline_curve, initialize_spline_curve):
-    return initialize_spline_curve(twice_differentiable_spline_curve)
+def initialised_twice_differentiable_spline_curve(twice_differentiable_spline_curve, initialise_spline_curve):
+    return initialise_spline_curve(twice_differentiable_spline_curve)
 
 
 @pytest.fixture
-def initialized_2D_spline_curve(spline_curve, initialize_2D_spline_curve):
-    return initialize_2D_spline_curve(spline_curve)
+def initialised_2D_spline_curve(spline_curve, initialise_2D_spline_curve):
+    return initialise_2D_spline_curve(spline_curve)
 
 
 @pytest.fixture
-def initialized_3D_spline_curve(spline_curve, initialize_3D_spline_curve):
-    return initialize_3D_spline_curve(spline_curve)
+def initialised_3D_spline_curve(spline_curve, initialise_3D_spline_curve):
+    return initialise_3D_spline_curve(spline_curve)
 
 
 @pytest.fixture
-def initialized_twice_differentiable_2D_spline_curve(twice_differentiable_spline_curve, initialize_2D_spline_curve):
-    return initialize_2D_spline_curve(twice_differentiable_spline_curve)
+def initialised_twice_differentiable_2D_spline_curve(twice_differentiable_spline_curve, initialise_2D_spline_curve):
+    return initialise_2D_spline_curve(twice_differentiable_spline_curve)
 
 
 @pytest.fixture
-def initialized_twice_differentiable_3D_spline_curve(twice_differentiable_spline_curve, initialize_3D_spline_curve):
-    return initialize_3D_spline_curve(twice_differentiable_spline_curve)
+def initialised_twice_differentiable_3D_spline_curve(twice_differentiable_spline_curve, initialise_3D_spline_curve):
+    return initialise_3D_spline_curve(twice_differentiable_spline_curve)
 
 
 @pytest.fixture
-def initialized_non_hermite_spline_curve(non_hermite_spline_curve, initialize_spline_curve):
-    return initialize_spline_curve(non_hermite_spline_curve)
+def initialised_non_hermite_spline_curve(non_hermite_spline_curve, initialise_spline_curve):
+    return initialise_spline_curve(non_hermite_spline_curve)
 
 
 @pytest.fixture
-def initialized_hermite_spline_curve(hermite_spline_curve, initialize_spline_curve):
-    return initialize_spline_curve(hermite_spline_curve)
+def initialised_hermite_spline_curve(hermite_spline_curve, initialise_spline_curve):
+    return initialise_spline_curve(hermite_spline_curve)
 
 
 @pytest.fixture
@@ -427,15 +444,18 @@ def boundary_condition(request):
     return request.param
 
 
-@pytest.fixture(params=[2, 5, 20, 50])
+@pytest.fixture(
+    params=[2, 5, 20, 50],
+    ids=["2 data points", "5 data points", "20 data points", "50 data points"],
+)
 def n_points(request):
     return request.param
 
 
 @pytest.fixture
-def points(codomain_dimensionality, n_points):
+def points(codomain_dimensionality, n_points, as_xp_array):
     rng = np.random.default_rng(seed=5544)
-    return rng.random((n_points, codomain_dimensionality))
+    return as_xp_array(rng.random((n_points, codomain_dimensionality)))
 
 
 @pytest.fixture(params=[False, True])
@@ -498,5 +518,5 @@ def single_value_t(request):
         np.array([1, 2, 3], dtype=int),
     ]
 )
-def array_t(request):
-    return request.param
+def array_t(request, as_xp_array):
+    return as_xp_array(request.param)

From ee3596631b1ab95de7b1d5d516ef41b54ff71c0e Mon Sep 17 00:00:00 2001
From: faymanns <florian.aymanns@epfl.ch>
Date: Sun, 11 Oct 2026 20:23:34 +0200
Subject: [PATCH 12/15] Test output of basis function calls

---
 src/splinebox/basis_functions.py | 106 +++++++++++++++++--------------
 tests/test_basis_functions.py    |  45 +++++++------
 2 files changed, 85 insertions(+), 66 deletions(-)

diff --git a/src/splinebox/basis_functions.py b/src/splinebox/basis_functions.py
index 052b716..a4cbe96 100644
--- a/src/splinebox/basis_functions.py
+++ b/src/splinebox/basis_functions.py
@@ -14,6 +14,7 @@
 import contextlib
 import inspect
 import math
+import numbers
 import sys
 import warnings
 
@@ -189,7 +190,7 @@ def __eq__(self, other):
             and other.support == self.support
         )
 
-    def __call__(self, t, derivative=0, jit=True):
+    def __call__(self, t, derivative=0, jit=False):
         """
         Evaluate the function at position(s) t. You can optionally
         pass derivative=1 or 2 to obtain the respective derivative(s).
@@ -222,7 +223,8 @@ def __call__(self, t, derivative=0, jit=True):
         """
         xp = array_api_compat.numpy if isinstance(t, (int, float)) else array_api_compat.array_namespace(t)
 
-        use_numba = jit and array_api_compat.is_numpy_namespace(xp)
+        # The xp functions can only accept arrays to for single value we use the numba version
+        use_numba = (jit and array_api_compat.is_numpy_namespace(xp)) or isinstance(t, numbers.Number)
 
         # Select the function that should be run
         if derivative == 0:
@@ -238,9 +240,10 @@ def __call__(self, t, derivative=0, jit=True):
         if use_numba:
             return func(t)
         elif jit and array_api_compat.is_torch_namespace(xp):
-            return _cached_compiled(
-                func, "torch", lambda fn: xp.compile(fullgraph=True, dynamic=True, mode="reduce-overhead")(fn)
-            )(t, xp)
+            with _torch_compile_recompile_limit():
+                return _cached_compiled(
+                    func, "torch", lambda fn: xp.compile(fullgraph=True, dynamic=True, mode="reduce-overhead")(fn)
+                )(t, xp)
         elif jit and array_api_compat.is_jax_namespace(xp):
             import jax
 
@@ -512,8 +515,8 @@ def _func_numba(t):  # pragma: no cover
     def _func_xp(t, xp):  # pragma: no cover
         return xp.where(
             xp.abs(t) < 1,
-            1 - xp.abs(t),
-            0,
+            1.0 - xp.abs(t),
+            0.0,
         )
 
     @staticmethod
@@ -533,7 +536,7 @@ def _derivative_1_xp(t, xp):  # pragma: no cover
         abs_t = xp.abs(t)
         result = -xp.sign(t)
         result = xp.where((t == 0) | (abs_t == 1), xp.nan, result)
-        return xp.where(abs_t > 1, 0, result)
+        return xp.where(abs_t > 1, 0.0, result)
 
     @staticmethod
     @numba.vectorize([numba.float64(numba.float64)], nopython=True, cache=True)
@@ -542,7 +545,7 @@ def _derivative_2_numba(t):
 
     @staticmethod
     def _derivative_2_xp(t, xp):
-        return xp.where((t == -1) | (t == 0) | (t == 1), xp.nan, 0)
+        return xp.where((t == -1) | (t == 0) | (t == 1), xp.nan, xp.asarray(0.0, dtype=t.dtype))
 
     @staticmethod
     def _filter_symmetric_numba(s):
@@ -669,7 +672,7 @@ def _func_xp(t, xp):  # pragma: no cover
                 xp.where(
                     (t > 0.5) & (t <= 1.5),
                     0.5 * (t**2) - 1.5 * t + 1.125,
-                    0,
+                    0.0,
                 ),
             ),
         )
@@ -685,7 +688,7 @@ def _derivative_1_xp(t, xp):  # pragma: no cover
                 xp.where(
                     (t > 0.5) & (t <= 1.5),
                     t - 1.5,
-                    0,
+                    0.0,
                 ),
             ),
         )
@@ -694,19 +697,19 @@ def _derivative_1_xp(t, xp):  # pragma: no cover
     def _derivative_2_xp(t, xp):  # pragma: no cover
         return xp.where(
             t < -1.5,
-            0,
+            0.0,
             xp.where(
                 (t > -1.5) & (t < -0.5),
-                1,
+                1.0,
                 xp.where(
                     (t > -0.5) & (t < 0.5),
-                    -2,
+                    -2.0,
                     xp.where(
                         (t > 0.5) & (t < 1.5),
-                        1,
+                        1.0,
                         xp.where(
                             t > 1.5,
-                            0,
+                            xp.asarray(0.0, dtype=t.dtype),
                             xp.nan,
                         ),
                     ),
@@ -819,7 +822,7 @@ def _func_xp(t, xp):  # pragma: no cover
             xp.where(
                 (abs_t >= 1) & (abs_t <= 2),
                 ((2 - abs_t) ** 3) / 6,
-                0,
+                0.0,
             ),
         )
 
@@ -837,7 +840,7 @@ def _derivative_1_xp(t, xp):  # pragma: no cover
                     xp.where(
                         (t >= -2) & (t <= -1),
                         0.5 * ((2 + t) ** 2),
-                        0,
+                        0.0,
                     ),
                 ),
             ),
@@ -857,7 +860,7 @@ def _derivative_2_xp(t, xp):  # pragma: no cover
                     xp.where(
                         (t >= -2) & (t <= -1),
                         2 + t,
-                        0,
+                        0.0,
                     ),
                 ),
             ),
@@ -1126,7 +1129,8 @@ def _func_xp(self, t, xp):  # pragma: no cover
         half_support = self.support / 2
         t = t + half_support
 
-        alpha = xp.asarray(xp.pi / self.M)
+        alpha = xp.pi / self.M
+        alpha = xp.asarray(alpha, dtype=xp.result_type(t, alpha), device=array_api_compat.device(t))
         L = 1 / (4 * xp.sin(alpha) ** 2)
 
         val = xp.where(
@@ -1138,7 +1142,7 @@ def _func_xp(self, t, xp):  # pragma: no cover
                 xp.where(
                     (t >= 2) & (t <= 3),
                     2 * xp.sin(alpha * (t - 3)) ** 2,
-                    0,
+                    0.0,
                 ),
             ),
         )
@@ -1148,7 +1152,8 @@ def _derivative_1_xp(self, t, xp):  # pragma: no cover
         half_support = self.support / 2
         t = t + half_support
 
-        alpha = xp.asarray(xp.pi / self.M)
+        alpha = xp.pi / self.M
+        alpha = xp.asarray(alpha, dtype=xp.result_type(t, alpha), device=array_api_compat.device(t))
         L = 1 / (4 * xp.sin(alpha) ** 2)
 
         val = xp.where(
@@ -1160,7 +1165,7 @@ def _derivative_1_xp(self, t, xp):  # pragma: no cover
                 xp.where(
                     (t > 2) & (t <= 3),
                     4 * alpha * xp.sin(alpha * (t - 3)) * xp.cos(alpha * (t - 3)),
-                    0,
+                    0.0,
                 ),
             ),
         )
@@ -1170,12 +1175,13 @@ def _derivative_2_xp(self, t, xp):  # pragma: no cover
         half_support = self.support / 2
         t = t + half_support
 
-        alpha = xp.asarray(xp.pi / self.M)
+        alpha = xp.pi / self.M
+        alpha = xp.asarray(alpha, dtype=xp.result_type(t, alpha), device=array_api_compat.device(t))
         L = 1 / (4 * xp.sin(alpha) ** 2)
 
         val = xp.where(
             t < 0,
-            0,
+            0.0,
             xp.where(
                 (t > 0) & (t < 1),
                 4 * alpha**2 * (xp.cos(alpha * t) ** 2 - xp.sin(alpha * t) ** 2),
@@ -1187,7 +1193,7 @@ def _derivative_2_xp(self, t, xp):  # pragma: no cover
                         4 * alpha**2 * (xp.cos(alpha * (t - 3)) ** 2 - xp.sin(alpha * (t - 3)) ** 2),
                         xp.where(
                             t > 3,
-                            0,
+                            0.0,
                             xp.nan,
                         ),
                     ),
@@ -1473,7 +1479,7 @@ def _func_xp(t, xp):  # pragma: no cover
             xp.where(
                 (abs_t > 1) & (abs_t <= 2),
                 (-1 / 2) * (abs_t**3) + (5 / 2) * (abs_t**2) - 4 * abs_t + 2,
-                0,
+                0.0,
             ),
         )
 
@@ -1491,7 +1497,7 @@ def _derivative_1_xp(t, xp):  # pragma: no cover
                     xp.where(
                         (t >= -2) & (t < -1),
                         1.5 * t * t + 5 * t + 4,
-                        0,
+                        0.0,
                     ),
                 ),
             ),
@@ -1517,7 +1523,7 @@ def _derivative_2_xp(t, xp):  # pragma: no cover
                             xp.where(
                                 (t > -2) & (t < -1),
                                 3 * t + 5,
-                                0,
+                                0.0,
                             ),
                         ),
                     ),
@@ -1681,7 +1687,7 @@ def h31_xp(t, xp):  # pragma: no cover
             xp.where(
                 (t < 0) & (t >= -1),
                 (1 - (2 * t)) * (t + 1) * (t + 1),
-                0,
+                0.0,
             ),
         )
 
@@ -1693,7 +1699,7 @@ def h32_xp(t, xp):  # pragma: no cover
             xp.where(
                 (t < 0) & (t >= -1),
                 t * (t + 1) * (t + 1),
-                0,
+                0.0,
             ),
         )
 
@@ -1708,7 +1714,7 @@ def h31prime_xp(t, xp):  # pragma: no cover
             xp.where(
                 (t < 0) & (t >= -1),
                 -6 * t * (t + 1),
-                0,
+                0.0,
             ),
         )
 
@@ -1720,7 +1726,7 @@ def h32prime_xp(t, xp):  # pragma: no cover
             xp.where(
                 (t < 0) & (t >= -1),
                 3 * t * t + 4 * t + 1,
-                0,
+                0.0,
             ),
         )
 
@@ -1735,7 +1741,7 @@ def h31primeprime_xp(t, xp):  # pragma: no cover
             xp.where(
                 (t < 0) & (t >= -1),
                 -12 * t - 6,
-                0,
+                0.0,
             ),
         )
 
@@ -1747,7 +1753,7 @@ def h32primeprime_xp(t, xp):  # pragma: no cover
             xp.where(
                 (t < 0) & (t >= -1),
                 6 * t + 4,
-                0,
+                0.0,
             ),
         )
 
@@ -2062,19 +2068,21 @@ def _func_xp(self, t, xp):  # pragma: no cover
     @staticmethod
     def _he31_xp(t, M, xp):  # pragma: no cover
         abs_t = xp.abs(t)
-        alpha = xp.asarray(xp.pi / M)
+        alpha = xp.pi / M
+        alpha = xp.asarray(alpha, dtype=xp.result_type(t, alpha), device=array_api_compat.device(t))
         denom = (alpha * xp.cos(alpha)) - xp.sin(alpha)
         num = (
             (0.5 * (2 * alpha * xp.cos(alpha) - xp.sin(alpha)))
             - (alpha * xp.cos(alpha) * abs_t)
             - (0.5 * xp.sin(alpha - (2 * alpha * abs_t)))
         )
-        return xp.where((abs_t >= 0) & (abs_t <= 1), num / denom, 0)
+        return xp.where((abs_t >= 0) & (abs_t <= 1), num / denom, 0.0)
 
     @staticmethod
     def _he32_xp(t, M, xp):  # pragma: no cover
         abs_t = xp.abs(t)
-        alpha = xp.asarray(xp.pi / M)
+        alpha = xp.pi / M
+        alpha = xp.asarray(alpha, dtype=xp.result_type(t, alpha), device=array_api_compat.device(t))
         denom = ((alpha * xp.cos(alpha)) - xp.sin(alpha)) * 8 * alpha * xp.sin(alpha)
         num = (
             -((2 * alpha * xp.cos(2 * alpha)) - xp.sin(2 * alpha))
@@ -2082,7 +2090,7 @@ def _he32_xp(t, M, xp):  # pragma: no cover
             - (2 * xp.sin(alpha) * xp.cos(2 * alpha * (abs_t - 0.5)))
             + (2 * alpha * xp.cos(2 * alpha * (abs_t - 1)))
         )
-        val = xp.where((abs_t >= 0) & (abs_t <= 1), num / denom, 0)
+        val = xp.where((abs_t >= 0) & (abs_t <= 1), num / denom, 0.0)
         return xp.where(t >= 0, val, -val)
 
     def _derivative_1_xp(self, t, xp):  # pragma: no cover
@@ -2091,23 +2099,25 @@ def _derivative_1_xp(self, t, xp):  # pragma: no cover
     @staticmethod
     def _he31prime_xp(t, M, xp):  # pragma: no cover
         abs_t = xp.abs(t)
-        alpha = xp.asarray(xp.pi / M)
+        alpha = xp.pi / M
+        alpha = xp.asarray(alpha, dtype=xp.result_type(t, alpha), device=array_api_compat.device(t))
         denom = (alpha * xp.cos(alpha)) - xp.sin(alpha)
         num = -(alpha * xp.cos(alpha)) + (alpha * xp.cos(alpha - (2 * alpha * abs_t)))
-        val = xp.where((abs_t >= 0) & (abs_t <= 1), num / denom, 0)
+        val = xp.where((abs_t >= 0) & (abs_t <= 1), num / denom, 0.0)
         return xp.where(t >= 0, val, -val)
 
     @staticmethod
     def _he32prime_xp(t, M, xp):  # pragma: no cover
         abs_t = xp.abs(t)
-        alpha = xp.asarray(xp.pi / M)
+        alpha = xp.pi / M
+        alpha = xp.asarray(alpha, dtype=xp.result_type(t, alpha), device=array_api_compat.device(t))
         denom = ((alpha * xp.cos(alpha)) - xp.sin(alpha)) * 8 * alpha * xp.sin(alpha)
         num = (
             -(4 * alpha * xp.sin(alpha) * xp.sin(alpha))
             + (4 * alpha * xp.sin(alpha) * xp.sin(2 * alpha * (abs_t - 0.5)))
             - (4 * alpha**2 * xp.sin(2 * alpha * (abs_t - 1)))
         )
-        return xp.where((abs_t >= 0) & (abs_t <= 1), num / denom, 0)
+        return xp.where((abs_t >= 0) & (abs_t <= 1), num / denom, 0.0)
 
     def _derivative_2_xp(self, t, xp):  # pragma: no cover
         return xp.stack([self._he31primeprime_xp(t, self.M, xp), self._he32primeprime_xp(t, self.M, xp)], axis=-1)
@@ -2115,20 +2125,22 @@ def _derivative_2_xp(self, t, xp):  # pragma: no cover
     @staticmethod
     def _he31primeprime_xp(t, M, xp):  # pragma: no cover
         abs_t = xp.abs(t)
-        alpha = xp.asarray(xp.pi / M)
+        alpha = xp.pi / M
+        alpha = xp.asarray(alpha, dtype=xp.result_type(t, alpha), device=array_api_compat.device(t))
         denom = (alpha * xp.cos(alpha)) - xp.sin(alpha)
         num = 2 * alpha**2 * xp.sin(alpha - (2 * alpha * abs_t))
-        return xp.where((abs_t >= 0) & (abs_t <= 1), num / denom, 0)
+        return xp.where((abs_t >= 0) & (abs_t <= 1), num / denom, 0.0)
 
     @staticmethod
     def _he32primeprime_xp(t, M, xp):  # pragma: no cover
         abs_t = xp.abs(t)
-        alpha = xp.asarray(xp.pi / M)
+        alpha = xp.pi / M
+        alpha = xp.asarray(alpha, dtype=xp.result_type(t, alpha), device=array_api_compat.device(t))
         denom = ((alpha * xp.cos(alpha)) - xp.sin(alpha)) * 8 * alpha * xp.sin(alpha)
         num = +(8 * alpha**2 * xp.sin(alpha) * xp.cos(2 * alpha * (abs_t - 0.5))) - (
             8 * alpha**3 * xp.cos(2 * alpha * (abs_t - 1))
         )
-        val = xp.where((abs_t >= 0) & (abs_t <= 1), num / denom, 0)
+        val = xp.where((abs_t >= 0) & (abs_t <= 1), num / denom, 0.0)
         return xp.where(t >= 0, val, -val)
 
     @staticmethod
diff --git a/tests/test_basis_functions.py b/tests/test_basis_functions.py
index e2b7451..da8e34d 100644
--- a/tests/test_basis_functions.py
+++ b/tests/test_basis_functions.py
@@ -1,5 +1,7 @@
+import collections
 import math
 
+import array_api_compat
 import numpy as np
 import pytest
 import splinebox.basis_functions
@@ -53,7 +55,7 @@ def test_refinement_mask(basis_function, is_locally_refinable, request):
     if isinstance(basis_function, splinebox.basis_functions.Exponential):
         # The refinement mask of the exponential basis function fails.
         # This might be because the value of M has to be changed after
-        # refinement. Since we don't have an imidiate need for the refinement
+        # refinement. Since we don't have an immediate need for the refinement
         # mask, we xfail the tests for now and will revisit this later.
         request.node.add_marker(pytest.mark.xfail)
     if is_locally_refinable(basis_function):
@@ -79,32 +81,37 @@ def test_refinement_mask(basis_function, is_locally_refinable, request):
             basis_function.refinement_mask()
 
 
-def test_call_output_type(basis_function, derivative, is_hermite_basis_function):
+def test_call_output_type(
+    basis_function, derivative, basis_function_call_positions, is_hermite_basis_function, as_xp_array, xp
+):
     """
     If a single value is provided, a single value should be returned.
-    Otherwise, a numpy array should be returned.
+    Otherwise, an array should be returned.
     """
-    for t in [
-        0.5,
-        np.array([0.5]),
-        np.array([-0.5, 0.0, 0.5]),
-        np.array([[-0.7, -0.2, 0.3], [-0.6, -0.1, 0.4], [-0.5, 0.0, 0.5], [-0.4, 0.1, 0.6]]),
-    ]:
+    t = basis_function_call_positions
 
-        vals = basis_function(t, derivative=derivative)
+    vals = basis_function(t, derivative=derivative)
+
+    if not isinstance(t, collections.abc.Iterable):
 
         if is_hermite_basis_function(basis_function):
-            assert isinstance(vals, np.ndarray)
-            assert vals.shape[-1] == 2
-            if isinstance(t, float):
-                assert vals.ndim == 1
-            else:
-                assert np.all(vals.shape[:-1] == t.shape)
+            assert vals.ndim == 1
+            assert vals.shape[0] == 2
+        else:
+            assert isinstance(vals, float)
+
+    else:
 
+        # Check that the output is an array from the xp namespace
+        assert array_api_compat.is_array_api_obj(vals)
+        assert array_api_compat.array_namespace(vals) is xp
+
+        if is_hermite_basis_function(basis_function):
+            assert np.all(vals.shape[:-1] == t.shape)
         else:
-            assert isinstance(vals, type(t))
-            if isinstance(t, np.ndarray):
-                assert np.all(vals.shape == t.shape)
+            assert np.all(vals.shape == t.shape)
+
+        assert xp.isdtype(vals.dtype, kind="real floating")
 
 
 def test_filter_output_type(basis_function, request):

From 9db3eaa3a9ede2570fef24c70ee7adbf3cb59aad Mon Sep 17 00:00:00 2001
From: faymanns <florian.aymanns@epfl.ch>
Date: Mon, 12 Oct 2026 09:29:24 +0200
Subject: [PATCH 13/15] Test that output of basis function calls has the right
 dtype

---
 src/splinebox/basis_functions.py |  2 +-
 tests/conftest.py                | 22 ++++++++++++++++++----
 tests/test_basis_functions.py    | 14 +++++++++++---
 3 files changed, 30 insertions(+), 8 deletions(-)

diff --git a/src/splinebox/basis_functions.py b/src/splinebox/basis_functions.py
index a4cbe96..2be04fa 100644
--- a/src/splinebox/basis_functions.py
+++ b/src/splinebox/basis_functions.py
@@ -1194,7 +1194,7 @@ def _derivative_2_xp(self, t, xp):  # pragma: no cover
                         xp.where(
                             t > 3,
                             0.0,
-                            xp.nan,
+                            xp.asarray(xp.nan, dtype=t.dtype),
                         ),
                     ),
                 ),
diff --git a/tests/conftest.py b/tests/conftest.py
index 364eb16..bb674e6 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -61,6 +61,19 @@ def device(xp_device):
     return xp_device[1]
 
 
+@pytest.fixture(params=["int", "float32", "float64"])
+def dtype(request, xp):
+    """
+    The dtype of the input array.
+
+    The output is expected to be of dtype float32 or
+    float64 if the input is of dtype float32 or float64 respectively and the
+    default floating point dtype of the array namespace if the input is of an
+    integer dtype.
+    """
+    return int if request.param == "int" else getattr(xp, request.param)
+
+
 @pytest.fixture
 def as_xp_array(xp, device):
     def _as_xp_array(obj, **kwargs):
@@ -259,15 +272,16 @@ def call_positions(request, as_xp_array):
     params=[
         0.5,
         1,
+        4,
         [0.5],
-        [0, 1],
-        [-0.5, 0.0, 0.5],
+        [0, 1, 4],
+        [-0.5, 0.0, 0.5, 4.1],
         [[-0.7, -0.2, 0.3], [-0.6, -0.1, 0.4], [-0.5, 0.0, 0.5], [-0.4, 0.1, 0.6]],
     ]
 )
-def basis_function_call_positions(request, as_xp_array):
+def basis_function_call_positions(request, as_xp_array, dtype):
     if isinstance(request.param, collections.abc.Iterable):
-        return as_xp_array(request.param)
+        return as_xp_array(request.param, dtype=dtype)
     return request.param
 
 
diff --git a/tests/test_basis_functions.py b/tests/test_basis_functions.py
index da8e34d..a851fe3 100644
--- a/tests/test_basis_functions.py
+++ b/tests/test_basis_functions.py
@@ -82,11 +82,15 @@ def test_refinement_mask(basis_function, is_locally_refinable, request):
 
 
 def test_call_output_type(
-    basis_function, derivative, basis_function_call_positions, is_hermite_basis_function, as_xp_array, xp
+    basis_function, derivative, basis_function_call_positions, dtype, is_hermite_basis_function, as_xp_array, xp
 ):
     """
     If a single value is provided, a single value should be returned.
     Otherwise, an array should be returned.
+
+    The value returned should be a float. For integer input the returned
+    value should have the default floating point dtype of the array namespace
+    while float32 and float64 input should be returned unchanged.
     """
     t = basis_function_call_positions
 
@@ -106,13 +110,17 @@ def test_call_output_type(
         assert array_api_compat.is_array_api_obj(vals)
         assert array_api_compat.array_namespace(vals) is xp
 
+        # Integer input should be cast to the default floating point dtype of
+        # the array namespace while float32 and float64 input keeps its dtype.
+        assert xp.isdtype(vals.dtype, kind="real floating")
+        expected_dtype = xp.asarray(1.0).dtype if dtype is int else dtype
+        assert vals.dtype == expected_dtype
+
         if is_hermite_basis_function(basis_function):
             assert np.all(vals.shape[:-1] == t.shape)
         else:
             assert np.all(vals.shape == t.shape)
 
-        assert xp.isdtype(vals.dtype, kind="real floating")
-
 
 def test_filter_output_type(basis_function, request):
 

From da9a8020d9096afb94ebffb30bf5ae9cce1ce8a1 Mon Sep 17 00:00:00 2001
From: faymanns <florian.aymanns@epfl.ch>
Date: Mon, 12 Oct 2026 10:34:43 +0200
Subject: [PATCH 14/15] Test that output of basis function filters has the
 right dtype

---
 src/splinebox/basis_functions.py | 12 ++++++++----
 tests/test_basis_functions.py    | 27 ++++++++++++++++++++++-----
 2 files changed, 30 insertions(+), 9 deletions(-)

diff --git a/src/splinebox/basis_functions.py b/src/splinebox/basis_functions.py
index 2be04fa..5f7f01c 100644
--- a/src/splinebox/basis_functions.py
+++ b/src/splinebox/basis_functions.py
@@ -269,7 +269,7 @@ def _derivative_2_numba(self, t):
     def _derivative_2_xp(self, t, xp):
         raise NotImplementedError(BasisFunction._unimplemented_message)
 
-    def filter_symmetric(self, s, jit=True):
+    def filter_symmetric(self, s, jit=False):
         """
         Returns a filtered version of the input s,
         used to convert knots into control points for an **open spline**.
@@ -316,7 +316,9 @@ def filter_symmetric(self, s, jit=True):
         """
         xp = array_api_compat.array_namespace(s)
 
-        s = xp.astype(s, float)
+        if not xp.isdtype(s.dtype, kind="real floating"):
+            # Convert to default floating type
+            s = xp.astype(s, xp.asarray(1.0).dtype)
 
         use_numba = jit and array_api_compat.is_numpy_namespace(xp)
 
@@ -351,7 +353,7 @@ def _filter_symmetric_numba(self, s):
     def _filter_symmetric_xp(self, s, xp):
         raise NotImplementedError(BasisFunction._unimplemented_message)
 
-    def filter_periodic(self, s, jit=True):
+    def filter_periodic(self, s, jit=False):
         """
         Returns a filtered version of the input s, used to convert
         knots into control points for a **closed spline**.
@@ -404,7 +406,9 @@ def filter_periodic(self, s, jit=True):
         """
         xp = array_api_compat.array_namespace(s)
 
-        s = xp.astype(s, float)
+        if not xp.isdtype(s.dtype, kind="real floating"):
+            # Convert to default floating type
+            s = xp.astype(s, xp.asarray(1.0).dtype)
 
         use_numba = jit and array_api_compat.is_numpy_namespace(xp)
 
diff --git a/tests/test_basis_functions.py b/tests/test_basis_functions.py
index a851fe3..09708b7 100644
--- a/tests/test_basis_functions.py
+++ b/tests/test_basis_functions.py
@@ -122,20 +122,37 @@ def test_call_output_type(
             assert np.all(vals.shape == t.shape)
 
 
-def test_filter_output_type(basis_function, request):
+def test_filter_output_type(basis_function, as_xp_array, dtype, xp, request):
 
     if isinstance(basis_function, splinebox.basis_functions.B2):
         # The filter_symmetric and filter_periodic are not implemented for B2
         request.node.add_marker(pytest.mark.xfail)
 
     for s in [np.array([1, 2, 3, 4, 5]), np.array([[1, 1], [2, 2], [3, 3], [4, 4], [5, 5]])]:
+        s = as_xp_array(s, dtype=dtype)
+        print(s.dtype)
         output = basis_function.filter_periodic(s)
-        assert isinstance(output, np.ndarray)
-        assert np.issubdtype(output.dtype, np.floating)
+        print(output.dtype)
+
+        assert array_api_compat.is_array_api_obj(output)
+        assert array_api_compat.array_namespace(output) is xp
+        # Integer input should be cast to the default floating point dtype of
+        # the array namespace while float32 and float64 input keeps its dtype.
+        assert xp.isdtype(output.dtype, kind="real floating")
+        expected_dtype = xp.asarray(1.0).dtype if dtype is int else dtype
+        print(expected_dtype)
+        assert output.dtype == expected_dtype
         assert np.all(output.shape == s.shape)
+
+        # Same test for the symmetric filter
         output = basis_function.filter_symmetric(s)
-        assert isinstance(output, np.ndarray)
-        assert np.issubdtype(output.dtype, np.floating)
+        assert array_api_compat.is_array_api_obj(output)
+        assert array_api_compat.array_namespace(output) is xp
+        # Integer input should be cast to the default floating point dtype of
+        # the array namespace while float32 and float64 input keeps its dtype.
+        assert xp.isdtype(output.dtype, kind="real floating")
+        expected_dtype = xp.asarray(1.0).dtype if dtype is int else dtype
+        assert output.dtype == expected_dtype
         assert np.all(output.shape == s.shape)
 
 

From 165d8f9d0c98d71d5a4b872d451f8bf854c023f7 Mon Sep 17 00:00:00 2001
From: faymanns <florian.aymanns@epfl.ch>
Date: Tue, 13 Oct 2026 17:31:21 +0200
Subject: [PATCH 15/15] Remove prints

---
 tests/test_basis_functions.py | 3 ---
 1 file changed, 3 deletions(-)

diff --git a/tests/test_basis_functions.py b/tests/test_basis_functions.py
index 09708b7..10fbbb7 100644
--- a/tests/test_basis_functions.py
+++ b/tests/test_basis_functions.py
@@ -130,9 +130,7 @@ def test_filter_output_type(basis_function, as_xp_array, dtype, xp, request):
 
     for s in [np.array([1, 2, 3, 4, 5]), np.array([[1, 1], [2, 2], [3, 3], [4, 4], [5, 5]])]:
         s = as_xp_array(s, dtype=dtype)
-        print(s.dtype)
         output = basis_function.filter_periodic(s)
-        print(output.dtype)
 
         assert array_api_compat.is_array_api_obj(output)
         assert array_api_compat.array_namespace(output) is xp
@@ -140,7 +138,6 @@ def test_filter_output_type(basis_function, as_xp_array, dtype, xp, request):
         # the array namespace while float32 and float64 input keeps its dtype.
         assert xp.isdtype(output.dtype, kind="real floating")
         expected_dtype = xp.asarray(1.0).dtype if dtype is int else dtype
-        print(expected_dtype)
         assert output.dtype == expected_dtype
         assert np.all(output.shape == s.shape)