-
Notifications
You must be signed in to change notification settings - Fork 9
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
feat!: support multiple backends #51
base: main
Are you sure you want to change the base?
Conversation
Install numpy by default, as it's lightweight and a dependency of tensorflow anyway. To use tensorflow, install phasespace as: pip install phasespace[tf]
Codecov Report
@@ 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
Continue to review full report at Codecov.
|
Sorry for the late reply, my previous reply got deleted underway.
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:
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 Feel free to change anything needed ;)
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 three that require the tensorflow probability? If a test requires TensorFlow, that is not a problem, right?
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
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 |
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! |
This reverts commit d2497677d910ee57a5ea0ef5c77fbd37ddf798f3.
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. |
Sure, many thanks for the code so far, that looks like what I had in mind so far, and yes, I agree |
Hi @redeboer I was wondering about the status of this? Do you think it's a good idea to finish this up soon? |
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 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. |
for more information, see https://pre-commit.ci
for more information, see https://pre-commit.ci
Follow-up to #46
Afaics, this setup allows one two switch between
numpy
andtensorflow
as backend forphasespace
, resulting in the same distributions. That said, there are some major issues:PHASESPACE_EAGER
variable, but I would prefer to switch the backend dynamically (smt likephasespace.set_backend("numpy")
).backend
sub-module works, but is emm, not as charming... It originates from the fact that not all TF operations are available intensorflow.experimental.numpy
, particularly whattensorflow.random
offers. So this module provides some kind of mapping (not dynamically :S) for TF functionsm thatphasespace
needs but that are not in the TF numpy API. Some problems with this approach:backend.random
'module' is tightly coupled tophasespace.random
: change one thing there, and each 'implementation' ofbackend.random
has to be adapted as well. In addition, one may wonder whetherphasespace.random.generate_uniform
shouldn't also be put underbackend
.backend.random
, as well asassert_equal
,shape
, etc follow TensorFlow's interface, whereas we want to move in the direction ofnumpy
(assumingtnp
does as well).tnp
) functions, you would have to add it to the 'mappings' for each backend.function
decorators: I had hoped to easily implement JAX as interface as well (and just usejax.jit
for those functions). This requires more changes though (if possible at all), because JAX can't handle custom classes (here:GenParticle
).phasespace
without the rather bulky TF, hence 3bd1f85. This is a breaking change though, as you would now have to usepip 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 :)