ENH: compile estimators entirely with JAX

[grayson-helmholz]

Closes #570

⚙️ Enhancements

On the JAX backend, ChiSquared and UnbinnedNLL are now JIT-compiled as a whole, instead of only JIT-compiling the wrapped function. The estimator body is built once in the constructor as a pure kernel over (parameters, *data) and wrapped in jax.jit, so the reduction, the normalization, and the function evaluation all fuse into a single compiled call that is reused across evaluations. A regression test asserts that a fit loop traces the kernel exactly once.

The analytic gradient() is now jax.jit(jax.grad(core)) over that same kernel, differentiated with respect to its parameters argument. It previously differentiated self.__call__, which forced a re-trace through the bound method on every call.

Data samples are converted to backend arrays once, in the constructor, so JIT-compiled evaluations no longer pay a host-to-device transfer per call.

❗ Behavioral changes

Estimators now import and initialize JAX when they are constructed rather than on the first gradient call, so a configure(jax_precision=...) call made before construction is respected. Estimator results are computed at the configured precision, where a fixed x64 flag was effectively assumed before.

Parameter values handed to an estimator are coerced to float or complex before evaluation. This keeps JIT input types stable — an int value would otherwise re-trace the kernel as soon as the optimizer turns it into a float — but it means integer parameter values no longer reach the wrapped function as integers.

Squash commit messages

* BEHAVIOR: coerce estimator parameters to float or complex
* BEHAVIOR: initialize JAX when constructing an estimator
* ENH: move estimator data to the backend arrays once

[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: 3b506e9	Previous: f5c1db9	Ratio
benchmarks/ampform.py::TestJPsiToGammaPiPi::test_fit[10000-jax]	0.6168525486949856 iter/sec (stddev: 0)	1.439081184425215 iter/sec (stddev: 0)	2.33

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