ENH: expose computational backend on function objects

.cspell.json

@@ -240,6 +240,7 @@
     "unbinned",
     "vectorize",
     "venv",
+    "vmap",
     "weisskopf",
     "wirtinger",
     "xcode",

docs/amplitude-analysis.ipynb

@@ -1907,7 +1907,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.13.13"
+   "version": "3.13.14"
   }
  },
  "nbformat": 4,

docs/usage.ipynb

@@ -805,7 +805,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.13.12"
+   "version": "3.13.14"
   }
  },
  "nbformat": 4,

docs/usage/basics.ipynb

@@ -1263,7 +1263,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.13.13"
+   "version": "3.13.14"
   }
  },
  "nbformat": 4,

docs/usage/binned-fit.ipynb

@@ -278,7 +278,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.13.12"
+   "version": "3.13.14"
   }
  },
  "nbformat": 4,

docs/usage/chi-squared.ipynb

@@ -231,7 +231,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.13.12"
+   "version": "3.13.14"
   }
  },
  "nbformat": 4,

docs/usage/unbinned-fit.ipynb

@@ -290,7 +290,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.13.12"
+   "version": "3.13.14"
   }
  },
  "nbformat": 4,

src/tensorwaves/data/transform.py

@@ -103,6 +103,6 @@ def from_sympy(
                 max_complexity=max_complexity,
             )
             functions[variable_name] = PositionalArgumentFunction(
-                function, argument_order
+                function, argument_order, backend
             )
         return cls(functions)

src/tensorwaves/estimator.py

@@ -75,6 +75,15 @@ def create_cached_function(
     return cached_function, cache_transformer
 
 
+def _determine_backend(function: ParametrizedFunction, backend: str | None) -> str:
+    if backend is not None:
+        return backend
+    function_backend = getattr(function, "backend", None)
+    if function_backend is None:
+        return "numpy"
+    return function_backend
+
+
 def gradient_creator(
     function: Callable[[Mapping[str, ParameterValue]], ParameterValue],
     backend: str,
@@ -120,7 +129,8 @@ class ChiSquared(Estimator):
             (unweighted). A common choice is :math:`w_i = 1/\sigma_i^2`, with
             :math:`\sigma_i` the uncertainty in each measured value of :math:`y_i`.
         backend: Computational backend with which to compute the sum
-            :math:`\sum_{i=1}^n`.
+            :math:`\sum_{i=1}^n`. By default, this is the backend of the
+            :code:`function`, if it exposes one (see `.BackendFunction`).
 
     .. seealso:: :doc:`/usage/chi-squared`
     """
@@ -131,8 +141,9 @@ def __init__(
         domain: DataSample,
         observed_values: np.ndarray,
         weights: np.ndarray | None = None,
-        backend: str = "numpy",
+        backend: str | None = None,
     ) -> None:
+        backend = _determine_backend(function, backend)
         self.__function = function
         self.__domain = domain
         self.__observed_values = observed_values
@@ -186,7 +197,8 @@ class UnbinnedNLL(Estimator):
         phsp_volume: Optional phase space volume :math:`V`, used in the
             normalization factor. Default: :math:`V=1`.
         backend: The computational back-end with which the sums and averages
-            should be computed.
+            should be computed. By default, this is the backend of the
+            :code:`function`, if it exposes one (see `.BackendFunction`).
 
     .. seealso:: :doc:`/usage/unbinned-fit`
     """
@@ -197,8 +209,9 @@ def __init__(
         data: DataSample,
         phsp: DataSample,
         phsp_volume: float = 1.0,
-        backend: str = "numpy",
+        backend: str | None = None,
     ) -> None:
+        backend = _determine_backend(function, backend)
         self.__data = dict(data)  # shallow copy
         self.__phsp = {k: v for k, v in phsp.items() if k != "weights"}
         self.__phsp_weights = phsp.get("weights")

src/tensorwaves/function/__init__.py

@@ -3,7 +3,7 @@
 from __future__ import annotations
 
 import inspect
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Protocol, runtime_checkable
 
 import attrs
 import numpy as np
@@ -20,6 +20,40 @@
     from collections.abc import Callable, Iterable, Mapping
 
 
+@runtime_checkable
+class BackendFunction(Protocol):
+    """A function object that exposes its backend-native kernel.
+
+    Classes like `PositionalArgumentFunction` and `ParametrizedBackendFunction` wrap a
+    pure function that takes positional argument arrays. This protocol gives access to
+    that kernel and the backend it was compiled for, so that backend-native
+    transformations (such as :code:`jax.jit`, :code:`jax.grad`, or :code:`jax.vmap`)
+    can be applied to it and estimators can determine which computational backend to
+    use.
+
+    >>> import sympy as sp
+    >>> from tensorwaves.function.sympy import create_function
+    >>> x, y = sp.symbols("x y")
+    >>> func = create_function(x**2 + y**2, backend="jax")
+    >>> func.backend
+    'jax'
+    >>> func.argument_order
+    ('x', 'y')
+    """
+
+    @property
+    def function(self) -> Callable[..., np.ndarray]:
+        """Backend-native function that takes positional arguments only."""
+
+    @property
+    def argument_order(self) -> tuple[str, ...]:
+        """Name of each positional argument, with data variables before parameters."""
+
+    @property
+    def backend(self) -> str | None:
+        """Name of the computational backend, if known."""
+
+
 def _all_str(
     _: PositionalArgumentFunction, __: attrs.Attribute, value: Iterable[str]
 ) -> None:
@@ -85,6 +119,8 @@ class PositionalArgumentFunction(Function[DataSample, np.ndarray]):
         converter=_to_tuple, validator=[_all_str, _all_unique]
     )
     """Ordered labels for each positional argument."""
+    backend: str | None = None
+    """Name of the computational backend that :attr:`function` was compiled for."""
 
     def __call__(self, data: DataSample) -> np.ndarray:
         args = [data[var_name] for var_name in self.argument_order]
@@ -102,8 +138,9 @@ def __init__(
         function: Callable[..., np.ndarray],
         argument_order: Iterable[str],
         parameters: Mapping[str, ParameterValue],
+        backend: str | None = None,
     ) -> None:
-        self.__function = PositionalArgumentFunction(function, argument_order)
+        self.__function = PositionalArgumentFunction(function, argument_order, backend)
         self.__parameters = dict(parameters)
 
     def __call__(
@@ -122,6 +159,10 @@ def function(self) -> Callable[..., np.ndarray]:
     def argument_order(self) -> tuple[str, ...]:
         return self.__function.argument_order
 
+    @property
+    def backend(self) -> str | None:
+        return self.__function.backend
+
     @property
     def parameters(self) -> dict[str, ParameterValue]:
         return dict(self.__parameters)
@@ -133,6 +174,7 @@ def with_parameters(
             function=self.function,
             argument_order=self.argument_order,
             parameters=self.__merge_parameters(parameters),
+            backend=self.backend,
         )
 
     def __merge_parameters(

src/tensorwaves/function/sympy/__init__.py

@@ -77,6 +77,7 @@ def create_function(
     return PositionalArgumentFunction(
         function=lambdified_function,
         argument_order=tuple(map(str, sorted_symbols)),
+        backend=backend,
     )
 
 
@@ -143,6 +144,7 @@ def create_parametrized_function(  # ruff:ignore[too-many-arguments]
         function=lambdified_function,
         argument_order=tuple(map(str, sorted_symbols)),
         parameters={str(symbol): value for symbol, value in parameters.items()},
+        backend=backend,
     )
 
 

tests/function/test_function.py

@@ -5,6 +5,7 @@
 import sympy as sp
 
 from tensorwaves.function import (
+    BackendFunction,
     ParametrizedBackendFunction,
     PositionalArgumentFunction,
     get_source_code,
@@ -58,6 +59,11 @@ def test_call(
     def test_function(self, function: ParametrizedBackendFunction):
         assert callable(function.function)
 
+    def test_backend(self, function: ParametrizedBackendFunction):
+        assert isinstance(function, BackendFunction)
+        assert function.backend == "numpy"
+        assert function.with_parameters({}).backend == "numpy"
+
     def test_call_with_parameters(self):
         initial_parameter_values = {"a": 1.0, "b": 2.0}
         func = ParametrizedBackendFunction(
@@ -87,6 +93,7 @@ def test_with_parameters(self):
         assert new_func is not func
         assert new_func.parameters == {"a": 2.0, "b": 2.0}
         assert new_func.function is func.function
+        assert func.backend is None
         assert func.parameters == initial_parameter_values
         data: DataSample = {"x": np.array([0.0, 1.0, 2.0])}
         np.testing.assert_array_equal(new_func(data), [2.0, 4.0, 6.0])

tests/test_estimator.py

@@ -42,6 +42,25 @@ def test_call(self, backend):
         )
         assert estimator({"a": 0, "b": 2}) == 2.5
 
+    @pytest.mark.parametrize("backend", ["jax", "numpy"])
+    def test_backend_inferred_from_function(self, backend):
+        x_data = {"x": np.array([0.0, 1.0, 2.0])}
+        y_data = np.array([0.0, 1.0, 2.0])
+        a, b, x = sp.symbols("a b x")
+        function = create_parametrized_function(
+            a + b * x,
+            parameters={a: 0.0, b: 1.0},
+            backend=backend,
+        )
+        estimator = ChiSquared(function, x_data, y_data)
+        if backend == "jax":
+            gradient = estimator.gradient({"a": 0.0, "b": 1.0})
+            assert pytest.approx(gradient["a"]) == 0.0
+            assert pytest.approx(gradient["b"]) == 0.0
+        else:
+            with pytest.raises(NotImplementedError):
+                estimator.gradient({"a": 0.0, "b": 1.0})
+
 
 def gaussian(mu_: float, sigma_: float) -> ParametrizedBackendFunction:
     x, mu, sigma = sp.symbols("x, mu, sigma")