feat!: support multiple backends

[redeboer]

• Implement numpy
• Implement JAX
• Switch back-ends dynamically (not just through environment variables)

---

Follow-up to #46

Afaics, this setup allows one two switch between numpy and tensorflow as backend for phasespace, resulting in the same distributions. That said, there are some major issues:

• I don't like that one can only switch between backends through an environment variable. I did this mostly to have static typing while developing and to mimic the existing PHASESPACE_EAGER variable, but I would prefer to switch the backend dynamically (smt like phasespace.set_backend("numpy")).
• Related to that: this backend sub-module works, but is emm, not as charming... It originates from the fact that not all TF operations are available in tensorflow.experimental.numpy, particularly what tensorflow.random offers. So this module provides some kind of mapping (not dynamically :S) for TF functionsm that phasespace needs but that are not in the TF numpy API. Some problems with this approach:
  • This backend.random 'module' is tightly coupled to phasespace.random: change one thing there, and each 'implementation' of backend.random has to be adapted as well. In addition, one may wonder whether phasespace.random.generate_uniform shouldn't also be put under backend.
  • Functions in backend.random, as well as assert_equal, shape, etc follow TensorFlow's interface, whereas we want to move in the direction of numpy (assuming tnp does as well).
  • This whole setup is not scalable: if you need some other TF (non-tnp) functions, you would have to add it to the 'mappings' for each backend.
• The backend with numpy is not fully tested: six tests are currently skipped because they requrie TF. The plots that are produced look okay though.
• With regard to the function decorators: I had hoped to easily implement JAX as interface as well (and just use jax.jit for those functions). This requires more changes though (if possible at all), because JAX can't handle custom classes (here: GenParticle).
• I would prefer if one can avoid installing phasespace without the rather bulky TF, hence 3bd1f85. This is a breaking change though, as you would now have to use pip install phasespace[tf] if you want to keep using TF as backend. Perhaps better to leave this for a follow-up PR, also so that that can be released separately. (Documentation would also have to be adapted.)

I'll leave this a draft PR, so we can think what to do about the above and other issues. Hopefully it's a step in the right direction :)

[codecov-commenter]

Codecov Report

Merging #51 (8413194) into master (b2d7c36) will decrease coverage by 4.69%.
The diff coverage is 68.42%.

@@            Coverage Diff             @@
##           master      #51      +/-   ##
==========================================
- Coverage   91.30%   86.61%   -4.70%     
==========================================
  Files           5        6       +1     
  Lines         345      381      +36     
  Branches       62       69       +7     
==========================================
+ Hits          315      330      +15     
- Misses         19       35      +16     
- Partials       11       16       +5     

Impacted Files	Coverage Δ
phasespace/backend/__init__.py	55.55% <55.55%> (ø)
phasespace/phasespace.py	89.10% <83.33%> (-0.52%)	⬇️
phasespace/random.py	86.66% <83.33%> (-13.34%)	⬇️
phasespace/__init__.py	100.00% <100.00%> (ø)
phasespace/backend/_tf_random.py	100.00% <100.00%> (ø)
phasespace/kinematics.py	98.03% <100.00%> (-0.04%)	⬇️

---

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update b2d7c36...8413194. Read the comment docs.

[jonas-eschle]

Sorry for the late reply, my previous reply got deleted underway.

• I don't like that one can only switch between backends through an environment variable. I did this mostly to have static typing while developing and to mimic the existing PHASESPACE_EAGER variable, but I would prefer to switch the backend dynamically (smt like phasespace.set_backend("numpy")).

Absolutely, I fully agree! I think it should be as you say. That's also how the eager actually works: the environment variable is just an additional way of enabling it but the functional way is there: tf.config.run_functions_eagerly. So to mimic it, we would indeed need a functional way as well, agree.

• Related to that: this backend sub-module works, but is emm, not as charming... It originates from the fact that not all TF operations are available in tensorflow.experimental.numpy, particularly what tensorflow.random offers. So this module provides some kind of mapping (not dynamically :S) for TF functionsm that phasespace needs but that are not in the TF numpy API.

Yes, first of all: there is no such thing as a backend switching algorithm in phasespace. It was designed to be easily implementable but it wasn't with the ability to. It was built on top of TF. So the backend could also be called tfextension or something currently. This means that anything concerning the backends, making the switchable is completely free to be changed, in fact even asked for; the whole structure for switching is not yet there. So I completely agree with your points, these are limitations by design and not made for a switchable backend

Feel free to change anything needed ;)

This whole setup is not scalable: if you need some other TF (non-tnp) functions, you would have to add it to the 'mappings' for each backend.

Isn't that an inherent problem when supporting multiple backends anyway? I don't see how one can possibly go around this, what exactly do you mean? Maybe an example?

The backend with numpy is not fully tested: six tests are currently skipped because they requrie TF. The plots that are produced look okay though.

The three that require the tensorflow probability? If a test requires TensorFlow, that is not a problem, right?

With regard to the function decorators: I had hoped to easily implement JAX as interface as well (and just use jax.jit for those functions). This requires more changes though (if possible at all), because JAX can't handle custom classes (here: GenParticle).

Yes, that is a current "disadvantages" I find of these JITing, but they are working towards making it easier to pass through composite objects. So I assume this will in the future be easier. And JAX will have something similar. A class is great if it serves basically as a namespace anyway.

But JAX also supports static arguments, you just have to specify which one it is with (the argnum or argname)[https://jax.readthedocs.io/en/latest/jax.html#jax.jit]. That is not perfect, but should work. And it means we need to wrap the wrappers maybe to have the same API with TF (that just ignores this argument)

I would prefer if one can avoid installing phasespace without the rather bulky TF, hence 3bd1f85. This is a breaking change though, as you would now have to use pip install phasespace[tf] if you want to keep using TF as backend. Perhaps better to leave this for a follow-up PR, also so that that can be released separately. (Documentation would also have to be adapted.)

Seems reasonable, but I agree, to put it into a longer term, e.g. for the 2.0 release. First to make it switchable and see, and once we're sure it's fine, we may can change the default. But for the moment, I would keep it installed.

On JAX: I would indeed favor to see this also in this PR, as we want to completely design here the backend strategy. Supporting one is one, two is two, but three is all in terms of backends. It's not a feature to support numpy but to have changeable backends and we will rather see the best way by designing for all (or three).

So please go ahead, that looks quite good so far and it seems we very well agree on the implementation and concepts

[jonas-eschle]

Hi, I was wondering if there are any news on this?

[redeboer]

Hi, I was wondering if there are any news on this?

Hi @jonas-eschle, sorry for not getting back to this. Have been focussing on the physics packages the past months and less so on the fitting/data generating.

I'm back next week and will have a look at it again. There seems to be quite some interest in switchable back-ends for the fitter side, so I definitely would like to move in that direction as well with data generation!

[redeboer]

I'm back next week and will have a look at it again.

Sorry, I have to get back to this later. What remains to be done is clear though, see updated description and your comments #51 (comment), and I'll get back to that as soon as I can. The challenging part will be JAX, but I agree that this (or some other third back-end) should be implemented in this PR as well.

[jonas-eschle]

Sure, many thanks for the code so far, that looks like what I had in mind so far, and yes, I agree

[jonas-eschle]

Hi @redeboer I was wondering about the status of this? Do you think it's a good idea to finish this up soon?

[redeboer]

Hi @jonas-eschle, thanks for pinging. I'm currently more focused on working out some physics parametrizations and less on the computational side in tensorwaves (now that that has become more independent from the physics). At some stage, I want to get this PR through, because we usually work with JAX/numpy as backend and don't want to have to pull in TF (through phasespace) on every install. It's not something that's high priority though.

If you prefer to clean up the PRs here, you can close this one for now and turn it into an issue. Then I'll backup this branch in a fork.