Can copies of arrays be avoided by XLA optimization?

[pfackeldey]

Dear JAX experts,

I am currently developing a fitting library in high energy physics for binned likelihood fits (dilax). These type of fits are typically used in physics analyses of particle physics experiments at CERN.

In these fits I need to evaluate a likelihood function, which is based on a large set of histograms (O(100 - 10^5)), and parameters that are adjusted during minimisation (e.g. experimental uncertainties). Each histogram has typically O(1-100) bins and the number of bins differ between histograms.
Currently I am just looping over histograms and apply functions that are based on previous mentioned parameters on them. However, I could in principle also concatenate multiple histograms together and then apply these functions. As far as I understand, I could benefit (significantly more) from vectorisation especially on GPUs; however, (afaik) jnp.concatenate will create copies of the inputs internally potentially leading to memory spikes, which is not the case when I just loop over these histograms.

My questions now boil down too:

1. Can XLA optimization potentially remove these copies? If yes, in which cases?
2. Is it in general better to jnp.concatenate for vectorisation, but potentially have memory spikes due to these copies, or is it better for loop over these different sized histograms?
3. Since I have a lot of these loops over different sets of histograms, I probably want to use jax.lax.scan to avoid large compile times (which I would get if I just use python for-loops as they are unrolled). Are there any other benefits of using jax.lax.scan, i.e. any optimizations which XLA can do compared to using a regular python for-loop?

Thank you very much in advance for you help!

Best, Peter

[jakevdp]

Hi - I don't believe that XLA will optimize-away copies in the case of concatenation. The problem is that input arrays will not generally be adjacent in memory, so the concatenated result cannot in general be represented using the original memory layout.

The best way to proceed would be to do your entire computation in a vectorized manner, e.g. by computing your histogram bins within a vmap if possible.

Hope that helps!

[pfackeldey]

Thank you very much for your reply! Yes, that makes sense to me :)

vmap can be applied in my case, but not in general, since I have thousands of small (irregular-sized) histograms. If some of them have the same number of bins I can merge them and use vmap (at the cost of some concatenation though).

Do you by chance have a short comment on my third question of the original post? This would clarify the benefits of jax.lax.scan to me:

Since I have a lot of these loops over different sets of histograms, I probably want to use jax.lax.scan to avoid large compile times (which I would get if I just use python for-loops as they are unrolled). Are there any other benefits of using jax.lax.scan, i.e. any optimizations which XLA can do compared to using a regular python for-loop?

Thank you very much for your help!

Best, Peter

[patrick-kidger]

On that last question -- it's usually the opposite. XLA finds it harder to optimize a scan than a Python for loop.