4 ms·
In my experience, the answer comes down to "does your code use classes liberally?" If no, you're just passing things between functions, then go ahead with Jax!
by polygamous_bat 3y ago
In my experience, the answer comes down to "does your code use classes liberally?"
If no, you're just passing things between functions, then go ahead with Jax! But converting larger codebases with classes is just significantly better with PyTorch even if they use different method names etc.
- patrickkidger 3y agoI'm going to disagree here! Classes and functional programming can go very well together, just don't expect to do in-place mutation. (I.e. OO-style programming.) You might like Equinox (https://github.com/patrick-kidger/equinox https://github.com/patrick-kidger/equinox ; 1.4k GitHub stars) which deliberately offers a very PyTorch-like feel for JAX. Regarding speed, I would strongly recommend JAX over PyTorch for SciComp. The XLA compiler seems to be much more effective for such use cases.