Speaker
Description
The JAX framework provides automatic differentiation, JIT compilation, vectorization, and multi-hardware acceleration well-suited for statistical inference in HEP. In this contribution, we present an ecosystem of interoperable tools that leverage the power of JAX, with a focus on everwillow, an inference tool agnostic to the underlying statistical model. At the modelling layer of this ecosystem, evermore provides an API for constructing binned likelihood models in JAX, while paramore builds on these primitives to support parametric, unbinned likelihoods. In parallel, pyhs3 leverages the HS3 workspace serialisation standard to import existing likelihoods into this ecosystem, including models originating from ROOT and RooFit. everwillow provides a common inference layer across these components. It requires only a likelihood function, whether defined via evermore or paramore, imported through pyhs3, or obtained from any other source that yields a JAX-compatible likelihood. This allows established workflows to benefit from JAX-based optimisation and acceleration. With everwillow, users can benefit from live fitting with real-time visualisation, fit checkpointing, hardware acceleration on multiple GPUs, novel computing opportunities such as vectorizing full fits, and analytical gradients. To show the power of this JAX ecosystem, we demonstrate cross-framework model combinations with custom parameter correlation schemes, joint inference, and statistical test evaluation. This outlines a path to modernising pythonic HEP statistical workflows while keeping compatibility with the ROOT ecosystem.