boost-histogram and traced computation (jax, TF and co)

[jonas-eschle]

Hi all, having used BH for a while and liking it, it's simple, fast and intuitive, I started wondering about integrating it within other workflows using jax, TensorFlow etc (specifically for binned fits in zfit). The peculiarity of it is that any Python code not written in TF, Jax... will not be rerun after the first execution. Or wrap the actual function with something like numpy_function that explicitly allows Python execution in a jitted function.

It seems somewhat similar to numba support #552 at first glance but may be different that jax should actually mimic the boost library histogram, as I am not quite sure whether, and how, that is actually possible to implement.

So the first question is maybe whether you have any thoughts on this? @HDembinski since you've also used jax, maybe you have an idea on good integration?

I see two (or three) possibilities:

• either we can stick close to BH and use as much functionality of it as possible. I've played around and basically implemented a _BaseHistogram with a TF/numpy-like implementation. With this, (and 2-3 changes in Histogram), I was able to get the functionality up for some test cases. This way would stick close to BH and maybe add to BH but has the disadvantage that we may would need to add (some) changes (I think minimal though) within BH. I see very well that this is rather not what is wanted for BH as it is just the bindings for the boost library. This has the problem that maybe a few things are incompatible in the API, such as *args, which are disliked by jax etc usually
• mimic BH interface as close as possible and have a completely independent implementation using TF, Jax... Maybe some functionality can be used from BH but on a best-effort basis. I think it may be the most realistic way but I don't quite like it obviously as it may has quite some code duplication
  (- maybe we can use directly BH by wrapping all the functions with numpy_function (something that adds a dynamic node inside the jitted function), but it seems no quite golden to me)

What are your general ideas on this?

[HDembinski]

Hi @mayou36 I think this is interesting and sorry that I somehow missed to comment, I hope @henryiii can also comment.

The peculiarity of it is that any Python code not written in TF, Jax... will not be rerun after the first execution.

I am not sure why this is a problem, the histogram itself should not change in whatever you do with TF and Jax. The histogram holds the data, while TF and JAX act on the model for the data. The data itself is static, you do not modify it in a fit. You only modify the predictions for the data as a function of the model parameters. So, in principle it should not be an issue at all that JAX or TF cannot interact directly with BH, it is sufficient if they can interface with the bin contents of BH, which can be converted to static numpy arrays and the bin edges, which can also be converted to static numpy arrays.

That being said, could you post a link or paste some code here that shows what kind of changes you think are required to implemented your first proposal?

[HDembinski]

To further elaborate on my point, we have a tutorial in Minuit that shows how to use JAX in fits of histogrammed data, see https://iminuit.readthedocs.io/en/latest/tutorial/automatic_differentiation.html. The histogramming here is done in plain numpy, it could have been done as easily with BH (I merely used Numpy to not add yet another dependency to my tutorials, which already feature a long list). JAX never touches the histogram, it only uses n and xe which are plain numpy arrays that you could also extract from BH.

[jonas-eschle]

Hi and sorry for the long silence, but I used your inputs to design the binned part from zfit.

The histogramming here is done in plain numpy, it could have been done as easily with BH (I merely used Numpy to not add yet another dependency to my tutorials, which already feature a long list). JAX never touches the histogram, it only uses n and xe which are plain numpy arrays that you could also extract from BH.

Indeed, but we could not pass in a histogram and/or use it's functionality in the jitted version.

That being said, could you post a link or paste some code here that shows what kind of changes you think are required to implemented your first proposal?

So I tried a bit, but I also got stuck somewhat, since there are minor other things that maybe don't work well inside a boost_histogram within a traced environment (bh works great of course! But it relies on knowing the values at execution time - a very reasonable assumption in general of course). So for the moment I ditched this effort and went with the following:

• created a BinnedData object within zfit that uses TensorFlow for all its computations. The interface of it is heavily inspired by bh and hist. It implements (basically) the UHI, it has "axes", counts, values, variances,... it can also be created using bh axis (Regular and Variable are currently supported).
• there are methods to go to and from hist.Hist method.
• The idea is to support a subset of what hist and bh offers, basically anything that makes sense inside a traced environment (say a .plot doesn't) and what is feasible to do with TF operations. We'll stick as close as possible to hist and bh, if you change something here in the design, we will try to reflect that and also change the BinnedData

Furthermore, just to make it easier:

• the loss function in zfit can directly take histograms (which, as a sidenote, are compatible with iminuit as they have the same interface as an iminuit cost function)
• so can basically anything that requires a BinnedData object (such as HistogramPDF)
• Any binned PDF can be converted to a histogram

I think this is the best approach, since most of it would anyway be a re-implementation to adhere to TF (or also JAX) jit tracing.

So we've reached a very high level of compatibility with bh & hist! Many thanks again for this two amazing libraries (and also the whole UHI story), it really pushed the development of binned fits in zfit.

P.S: the PR is still open in zfit, but should soon be out for a beta release