From f5b482597110e4e69f4e13e61ebc18bcecb49d18 Mon Sep 17 00:00:00 2001 From: Remco de Boer <29308176+redeboer@users.noreply.github.com> Date: Wed, 5 Aug 2026 22:01:17 +0200 Subject: [PATCH] ENH: JIT-compile estimators entirely with JAX --- .cspell.json | 1 + src/tensorwaves/estimator.py | 180 +++++++++++++++++++++++++++-------- tests/test_config.py | 24 +++++ tests/test_estimator.py | 22 +++++ 4 files changed, 188 insertions(+), 39 deletions(-) diff --git a/.cspell.json b/.cspell.json index c9134483..3dd99ec6 100644 --- a/.cspell.json +++ b/.cspell.json @@ -186,6 +186,7 @@ "version": "0.2", "words": [ "analyticity", + "argnums", "backends", "blatt", "bottomness", diff --git a/src/tensorwaves/estimator.py b/src/tensorwaves/estimator.py index 052d0c8a..07470c3d 100644 --- a/src/tensorwaves/estimator.py +++ b/src/tensorwaves/estimator.py @@ -84,24 +84,45 @@ def _determine_backend(function: ParametrizedFunction, backend: str | None) -> s 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 @@ -114,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. @@ -132,6 +193,9 @@ class ChiSquared(Estimator): :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` """ @@ -144,27 +208,45 @@ def __init__( backend: str | None = None, ) -> None: backend = _determine_backend(function, backend) - self.__function = function - self.__domain = domain - self.__observed_values = observed_values + 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): @@ -200,6 +282,9 @@ class UnbinnedNLL(Estimator): 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` """ @@ -212,28 +297,45 @@ def __init__( 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") - 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 + 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, + ) diff --git a/tests/test_config.py b/tests/test_config.py index e6246811..84d96e4a 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -101,6 +101,30 @@ def test_configure_before_creating_arrays(precision: str): assert _run(code) == precision +@pytest.mark.parametrize("precision", ["float32", "float64"]) +def test_jax_estimator_respects_precision(precision: str): + """Estimators initialize JAX through :func:`.configure`, not with a fixed x64 flag.""" + code = f""" +import numpy as np +import sympy as sp +from tensorwaves import configure +from tensorwaves.estimator import ChiSquared +from tensorwaves.function.sympy import create_parametrized_function +configure(jax_precision={precision!r}) +x, a = sp.symbols("x a") +function = create_parametrized_function(a * x, {{a: 1.0}}, backend="jax") +estimator = ChiSquared( + function, + domain={{"x": np.linspace(0, 1, num=10)}}, + observed_values=np.zeros(10), +) +import jax +print(jax.config.x64_enabled, estimator({{"a": 1.0}}).dtype.name) +""" + x64_enabled = precision == "float64" + assert _run(code) == f"{x64_enabled} {precision}" + + @pytest.mark.parametrize( argnames=("precision", "expected"), argvalues=[ diff --git a/tests/test_estimator.py b/tests/test_estimator.py index bf4ec966..1b44ddb9 100644 --- a/tests/test_estimator.py +++ b/tests/test_estimator.py @@ -42,6 +42,28 @@ def test_call(self, backend): ) assert estimator({"a": 0, "b": 2}) == 2.5 + def test_jit_compiled_once(self): + trace_count = 0 + + def linear(a, b, x): + nonlocal trace_count + trace_count += 1 + return a + b * x + + function = ParametrizedBackendFunction( + linear, + argument_order=("a", "b", "x"), + parameters={"a": 0.0, "b": 1.0}, + backend="jax", + ) + x_data = {"x": np.array([0.0, 1.0, 2.0])} + y_data = np.array([0.0, 2.0, 4.0]) + estimator = ChiSquared(function, x_data, y_data) + for b in [1.0, 2.0, 3.0]: + estimator({"a": 0.0, "b": b}) + assert trace_count == 1, "estimator was re-traced during evaluation" + assert estimator({"a": 0.0, "b": 2.0}) == 0.0 + @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])}