FEAT: expose backend on function objects

[grayson-helmholz]

Closes #569

✨ New features

New BackendFunction protocol in tensorwaves.function, which exposes the backend-native kernel of a function object together with the backend it was lambdified for. This makes it possible to apply backend-native transformations such as jax.jit, jax.grad, or jax.vmap to the wrapped kernel, and lets consumers discover which backend a function computes on. The protocol is runtime_checkable, so isinstance(function, BackendFunction) works.

PositionalArgumentFunction and ParametrizedBackendFunction now carry a backend attribute, which create_function() and create_parametrized_function() set from their backend argument. with_parameters() carries it over to the new function.

⚙️ Enhancements

The backend argument of ChiSquared and UnbinnedNLL now defaults to None, in which case the estimator infers the backend from the function it wraps. Passing backend= explicitly still overrides this.

❗ Behavioral changes

An estimator that wraps a function created with a non-NumPy backend now computes its sums and averages on that backend, where it previously always fell back to NumPy. In particular, ChiSquared(jax_function, ...).gradient(...) now returns an analytic JAX gradient instead of raising NotImplementedError. Estimators wrapping a function that exposes no backend keep computing on NumPy.

Squash commit messages

* BEHAVIOR: compute estimators on the function backend
* ENH: infer estimator backend from the function

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[redeboer]

⚠️ Performance Alert ⚠️

Possible performance regression was detected for benchmark 'TensorWaves benchmark results'.
Benchmark result of this commit is worse than the previous benchmark result exceeding threshold 2.

Benchmark suite	Current: e313133	Previous: dadc908	Ratio
benchmarks/expression.py::test_fit[1000-ScipyMinimizer-numpy]	7.087714707964495 iter/sec (stddev: 0.0005488246113790687)	15.653084369591935 iter/sec (stddev: 0.0008327728665216672)	2.21
benchmarks/expression.py::test_fit[1000-ScipyMinimizer-numba]	7.100748203441732 iter/sec (stddev: 0.00037290296655397926)	15.57675266668666 iter/sec (stddev: 0.00019544090300604612)	2.19
benchmarks/expression.py::test_fit[1000-ScipyMinimizer-tf]	0.6362030851825524 iter/sec (stddev: 0.008273548261204305)	1.4336772024447277 iter/sec (stddev: 0.002264740294580996)	2.25
benchmarks/unbinned_nll.py::test_unbinned_nll_normalization_formula[original-numpy]	28.910685999949724 iter/sec (stddev: 0.000178942339813722)	82.89554185733245 iter/sec (stddev: 0.00047803677168886747)	2.87
benchmarks/unbinned_nll.py::test_unbinned_nll_normalization_formula[optimized-numpy]	37.30708326842129 iter/sec (stddev: 0.00010286166659039556)	128.51191203721672 iter/sec (stddev: 0.00026220735978772406)	3.44
benchmarks/unbinned_nll.py::test_unbinned_nll_estimator[numpy]	124.83132538681305 iter/sec (stddev: 0.00017531715573278447)	345.5347754466312 iter/sec (stddev: 0.00009775185162269724)	2.77
benchmarks/unbinned_nll.py::test_unbinned_nll_estimator[numba]	99.52917181454815 iter/sec (stddev: 0.0001208770765365526)	547.3808519863164 iter/sec (stddev: 0.00008040495616927644)	5.50

This comment was automatically generated by workflow using github-action-benchmark.