Proof-of-concept idea: pyhf models as JAX PyTrees

[phinate]

This is a quickly thrown together draft of a proof-of-concept that I think would be both cool to see, and possible to prototype in a short timespan, mainly to see what's possible!

I don't imagine this kind of technical prototype as a replacement for pyhf, but I think that there could be lessons learned by working on this functionality that could be upstreamed to the current codebase. Hope these ideas are useful :)

What?

PyTrees are JAX's native representation of structured data. Many native Python types are already PyTrees: dict, list, NamedTuple etc. being some examples. A class, by default, is not a PyTree, but thanks to libraries like equinox and simple-pytree, one can just inherit from the relevant class provided by these modules, and suddenly your class is also a PyTree (for example, pyhf.Model!).

Now, why are PyTrees useful? The main reason is that they're compatible with JAX function transforms, meaning that functions that take in PyTrees as arguments can be used with jax.jit, jax.vmap, jax.grad etc., which opens up a whole range of cool stuff. I'll try to summarize exactly how this could be useful for pyhf in the following sections.

The main idea here is basically: if a module is designed with the right JAX building blocks, you can arbitrarily use any of its cool features on any of your programs! (this is the design principle behind modules like diffrax, as pointed out to me by the author).

Why?

vectorizable limits/hypotest/anything!

Vectorization/batching becomes super easy. E.g. for hypotest:

# what we normally do
pyhf.infer.hypotest(test_mu, data, model)

# get CLs for many values of mu
jax.vmap(pyhf.infer.hypotest, in_axis=(0, None, None)(test_mu_list, data, model)

# or maybe for many models
jax.vmap(pyhf.infer.hypotest, in_axis=(None, None, 0)(test_mu, data, model_list)

# now jit compile the whole thing
jax.jit(jax.vmap(pyhf.infer.hypotest, in_axis=(None, None, 0)(test_mu, data, model_list))

# and even take the gradient with respect to mu!
jax.grad(jax.jit(jax.vmap(pyhf.infer.hypotest, in_axis=(None, None, 0)))(test_mu, data, model_list)

This should be particularly nice for limit scanning -- this is currently a Python for-loop in the code:

pyhf/src/pyhf/infer/intervals/upper_limits.py

Lines 189 to 192 in 5161749

	results = [
	hypotest(mu, data, model, return_expected_set=True, **hypotest_kwargs)
	for mu in scan
	]

Access to jax.jit

All workflows that involve PyTrees get access to just-in-time compilation, which, like all JAX transforms, can be used as a function transformation or as a decorator:

def my_pyhf_workflow(spec):
    model = pyhf.Model(spec)
    ...do stuff
    return five_sigma_limits
    
jitted_workflow = jax.jit(my_pyhf_workflow)
res = jitted_workflow(spec)

# or just this 
@jax.jit
def my_pyhf_workflow(spec):
    model = pyhf.Model(spec)
    ...do stuff
    return five_sigma_limits

As a very small anecdotal benchmark, I messaged @matthewfeickert a couple years back claiming that logpdf could already be made 10x faster if the method is wrapped with jax.jit (I wonder if he remembers this :p)

True end-to-end differentiability

If we have PyTrees, all the differentiable model building stuff is super easy, resolving a number of issues and PRs I've made (e.g. #882, #1912). I found this out via experimenting to resolve this jaxopt issue, which has basically stopped relaxed from working on modern JAX versions.

See this conversation to see how easy things are with PyTrees!

How?

This project feels much more realistic after looking through @pfackeldey's work on dilax, which uses all the main building blocks I've discussed here -- I even managed to differentiate over model construction without him knowing about neos! My hunch for a realistic path forward is to use the same design principles, but to build things up in a more pyhf-like way such that the resulting API would look more or less the same. It's a great launchpad for this effort, if it is of interest!

Why haven't I done this yet?

I've toyed with this for a little while, but I don't have enough fundamental understanding of the best way to compose all the different parts of histfactory (mostly ignorance around systematics handling + computation). So with a little advice and a couple of people looking at this (e.g. as a hackathon project for PyHEP.dev, which I think most mentioned here should be attending), we could probably make this happen!

Oh, and I'm also technically not in the field right now. But I have 10% bandwidth that I'm allowed to use to write things like this ;)

My final, selfish reason for this is this is the make-or-break ingredient for things like neos being possible for anyone that wants to try it in the future with a HiFa-based likelihood (i've already had a couple people talk to me about that).

So yeah -- hope this is useful, and maybe it turns into something cool!

[nsmith-]

I like pytrees a lot, but in a PoC I was working on two years ago, I found that the jit time of using pytrees of many small arrays was very expensive. I found that it could take up to hours to jit a value_and_grad for a fit of a few hundred channels with only O(10) bins each (a CMS VHbb analysis). I had to progressively merge bins/processes/systematics into fewer contiguous vectors before the jit time was acceptable. e.g. some of the commits shown here https://github.com/nsmith-/jaxfit/commits/roofit
Granted, this was two years ago, so it would be very interesting to re-investigate the scaling properties of jax.jit + pytrees.

[nsmith-]

I can't find the document now, but at the time the recommendation was to check the length of the jax.make_jaxpr output as a proxy for jit time while developing/optimizing. Basically one wants to have the most compact intermediate repr so that XLA compilation happens quickly.

[pfackeldey]

Hi @nsmith- ! thanks for sharing your experience with long compile times. Have you by chance found out which part introduced this? This can happen very fast eg if one uses python loops instead of the corresponding jax primitives, such as jax.lax.fori_loop, jax.lax.while_loop, ... as python loops will just unrolled and create a lot of MHLO lines for XLA (as far as I understand). I could imagine the optimizer implementation to be responsible for this and wonder if you could see the compile times primarily arise from there. Definitely one needs a batchable optimizer (see eg jaxopt) otherwise many benefits @phinate listed above are not possible, I wonder how jit friendly (regarding compile times) these (jaxopt) are implemented?

[nsmith-]

The compilation was just for value and gradient, not including the minimizer code (which indeed does use jax.lax.while_loop). The issue with using lax loops in the likelihood function is that the arrays are irregular sizes. I solved this in the vectorization steps by padding and then creating a new axis that could then be summed over with the usual jax.numpy functions. This could also be a feature, as one could use ragged array techniques to concatenate the irregular arrays into something more efficient to compute.

[pfackeldey]

That's interesting! I have to admit that I do not see where ragged arrays come into play exactly. In my implementation (dilax) I have not encountered a place where I need to use ragged arrays yet. Is that maybe a choice of the initial data structure? E.g. In my case I have as data structure only the bins, and then apply functions on all (or parts) of these bin to modify them according to a nuisance parameter.