ENH: compile estimators entirely with JAX

.cspell.json

@@ -186,6 +186,7 @@
   "version": "0.2",
   "words": [
     "analyticity",
+    "argnums",
     "backends",
     "blatt",
     "bottomness",
@@ -240,6 +241,7 @@
     "unbinned",
     "vectorize",
     "venv",
+    "vmap",
     "weisskopf",
     "wirtinger",
     "xcode",

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,24 +75,54 @@ 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 _coerce_parameter_types(
+    parameters: Mapping[str, ParameterValue],
+) -> dict[str, ParameterValue]:
+    # normalize to float/complex so that JIT compilers see stable input types
+    # (an int value would otherwise trigger a re-trace once it becomes a float)
+    return {
+        name: complex(value) if isinstance(value, complex) else float(value)
+        for name, value in parameters.items()
+    }
+
+
+def _import_jax():  # ruff: ignore[missing-return-type-private-function]
+    try:
+        return _initialize_jax()
+    except ImportError:  # pragma: no cover
+        raise_missing_module_error("jax", extras_require="jax")
+
+
+def _conjugate_complex_gradient(
+    gradient: Mapping[str, ParameterValue],
+) -> dict[str, ParameterValue]:
+    # jax.grad() returns the conjugated Wirtinger derivative ∂f/∂x - i∂f/∂y
+    # for complex-valued parameters, so conjugate to get a complex number
+    # whose real and imaginary parts are (∂f/∂x, ∂f/∂y)
+    return {name: value.conjugate() for name, value in gradient.items()}
+
+
 def gradient_creator(
     function: Callable[[Mapping[str, ParameterValue]], ParameterValue],
     backend: str,
 ) -> Callable[[Mapping[str, ParameterValue]], dict[str, ParameterValue]]:
     if backend == "jax":
-        try:
-            jax = _initialize_jax()
-        except ImportError:  # pragma: no cover
-            raise_missing_module_error("jax", extras_require="jax")
+        jax = _import_jax()
         gradient = jax.grad(function)
 
         def conjugated_gradient(
             parameters: Mapping[str, ParameterValue],
         ) -> dict[str, ParameterValue]:
-            # jax.grad() returns the conjugated Wirtinger derivative ∂f/∂x - i∂f/∂y
-            # for complex-valued parameters, so conjugate to get a complex number
-            # whose real and imaginary parts are (∂f/∂x, ∂f/∂y)
-            return {k: v.conjugate() for k, v in gradient(parameters).items()}
+            return _conjugate_complex_gradient(gradient(parameters))
 
         return conjugated_gradient
 
@@ -105,6 +135,46 @@ def raise_gradient_not_implemented(
     return raise_gradient_not_implemented
 
 
+def _jit_estimator_core(core: Callable, backend: str) -> Callable:
+    if backend == "jax":
+        jax = _import_jax()
+        return jax.jit(core)
+    return core
+
+
+def _convert_arrays_to_backend(data: DataSample, backend: str) -> DataSample:
+    # move data arrays to the device once, so that JIT-compiled estimator calls
+    # do not pay a host-to-device transfer on every evaluation
+    if backend == "jax":
+        jax = _import_jax()
+        return {key: jax.numpy.asarray(array) for key, array in data.items()}
+    return data
+
+
+def _create_core_gradient(core: Callable, backend: str) -> Callable:
+    """Create a JIT-compiled gradient of an estimator core, w.r.t. its parameters."""
+    if backend == "jax":
+        jax = _import_jax()
+        raw_gradient = jax.jit(jax.grad(core, argnums=0))
+
+        def gradient(
+            parameters: Mapping[str, ParameterValue],
+            *data_args: DataSample | np.ndarray | None,
+        ) -> dict[str, ParameterValue]:
+            return _conjugate_complex_gradient(raw_gradient(parameters, *data_args))
+
+        return gradient
+
+    def raise_gradient_not_implemented(
+        parameters: Mapping[str, ParameterValue],
+        *data_args: DataSample | np.ndarray | None,
+    ) -> dict[str, ParameterValue]:
+        msg = f"Gradient not implemented for back-end {backend}."
+        raise NotImplementedError(msg)
+
+    return raise_gradient_not_implemented
+
+
 class ChiSquared(Estimator):
     r"""Chi-squared test estimator.
 
@@ -120,7 +190,11 @@ 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`).
+
+    On the JAX backend, the full estimator and its analytic :meth:`gradient` are
+    JIT-compiled once and cached over all further evaluations.
 
     .. seealso:: :doc:`/usage/chi-squared`
     """
@@ -131,29 +205,48 @@ def __init__(
         domain: DataSample,
         observed_values: np.ndarray,
         weights: np.ndarray | None = None,
-        backend: str = "numpy",
+        backend: str | None = None,
     ) -> None:
-        self.__function = function
-        self.__domain = domain
-        self.__observed_values = observed_values
+        backend = _determine_backend(function, backend)
+        self.__domain = _convert_arrays_to_backend(domain, backend)
         if weights is None:
             ones = find_function("ones", backend)
-            self.__weights = ones(len(self.__observed_values))
-        else:
-            self.__weights = weights
+            weights = ones(len(observed_values))
+        converted = _convert_arrays_to_backend(
+            {"observed_values": observed_values, "weights": weights}, backend
+        )
+        self.__observed_values = converted["observed_values"]
+        self.__weights = converted["weights"]
+        sum_function = find_function("sum", backend)
+
+        def estimator(
+            parameters: Mapping[str, ParameterValue],
+            domain: DataSample,
+            observed_values: np.ndarray,
+            weights: np.ndarray,
+        ) -> float:
+            computed_values = function(domain, parameters)
+            chi_squared = weights * (computed_values - observed_values) ** 2
+            return sum_function(chi_squared)
 
-        self.__gradient = gradient_creator(self.__call__, backend)
-        self.__sum = find_function("sum", backend)
+        self.__estimator = _jit_estimator_core(estimator, backend)
+        self.__gradient = _create_core_gradient(estimator, backend)
 
     def __call__(self, parameters: Mapping[str, ParameterValue]) -> float:
-        computed_values = self.__function(self.__domain, parameters)
-        chi_squared = self.__weights * (computed_values - self.__observed_values) ** 2
-        return self.__sum(chi_squared)
+        return self.__estimator(*self.__estimator_args(parameters))
 
     def gradient(
         self, parameters: Mapping[str, ParameterValue]
     ) -> dict[str, ParameterValue]:
-        return self.__gradient(parameters)
+        return self.__gradient(*self.__estimator_args(parameters))
+
+    def __estimator_args(self, parameters: Mapping[str, ParameterValue]) -> tuple:
+        return (
+            _coerce_parameter_types(parameters),
+            self.__domain,
+            self.__observed_values,
+            self.__weights,
+        )
 
 
 class UnbinnedNLL(Estimator):
@@ -186,7 +279,11 @@ 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`).
+
+    On the JAX backend, the full estimator and its analytic :meth:`gradient` are
+    JIT-compiled once and cached over all further evaluations.
 
     .. seealso:: :doc:`/usage/unbinned-fit`
     """
@@ -197,30 +294,48 @@ def __init__(
         data: DataSample,
         phsp: DataSample,
         phsp_volume: float = 1.0,
-        backend: str = "numpy",
+        backend: str | None = None,
     ) -> None:
-        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")
-        self.__function = function
-        self.__gradient = gradient_creator(self.__call__, backend)
-
-        self.__mean = find_function("mean", backend)
-        self.__sum = find_function("sum", backend)
-        self.__log = find_function("log", backend)
-
-        self.__phsp_volume = phsp_volume
+        backend = _determine_backend(function, backend)
+        self.__data = _convert_arrays_to_backend(dict(data), backend)
+        converted_phsp = _convert_arrays_to_backend(dict(phsp), backend)
+        self.__phsp = {k: v for k, v in converted_phsp.items() if k != "weights"}
+        self.__phsp_weights = converted_phsp.get("weights")
+        mean_function = find_function("mean", backend)
+        sum_function = find_function("sum", backend)
+        log_function = find_function("log", backend)
+
+        def estimator(
+            parameters: Mapping[str, ParameterValue],
+            data: DataSample,
+            phsp: DataSample,
+            phsp_weights: np.ndarray | None,
+        ) -> float:
+            bare_intensities = function(data, parameters)
+            phsp_intensities = function(phsp, parameters)
+            if phsp_weights is not None:
+                phsp_intensities *= phsp_weights
+            normalization_integral = phsp_volume * mean_function(phsp_intensities)
+            log_normalization = len(bare_intensities) * log_function(
+                normalization_integral
+            )
+            return log_normalization - sum_function(log_function(bare_intensities))
+
+        self.__estimator = _jit_estimator_core(estimator, backend)
+        self.__gradient = _create_core_gradient(estimator, backend)
 
     def __call__(self, parameters: Mapping[str, ParameterValue]) -> float:
-        data_intensities = self.__function(self.__data, parameters)
-        phsp_intensities = self.__function(self.__phsp, parameters)
-        if self.__phsp_weights is not None:
-            phsp_intensities *= self.__phsp_weights
-        normalization_integral = self.__phsp_volume * self.__mean(phsp_intensities)
-        log_normalization = len(data_intensities) * self.__log(normalization_integral)
-        return log_normalization - self.__sum(self.__log(data_intensities))
+        return self.__estimator(*self.__estimator_args(parameters))
 
     def gradient(
         self, parameters: Mapping[str, ParameterValue]
     ) -> dict[str, ParameterValue]:
-        return self.__gradient(parameters)
+        return self.__gradient(*self.__estimator_args(parameters))
+
+    def __estimator_args(self, parameters: Mapping[str, ParameterValue]) -> tuple:
+        return (
+            _coerce_parameter_types(parameters),
+            self.__data,
+            self.__phsp,
+            self.__phsp_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,
     )