3 ms·

It seems like the ecosystem is still dominated by PyTorch, is Jax supposed to be a competitor? Any signs of Jax taking over PyTorch anytime soon? Is it perhaps

It seems like the ecosystem is still dominated by PyTorch, is Jax supposed to be a competitor? Any signs of Jax taking over PyTorch anytime soon? Is it perhaps too early for its time? Or is there a critical flaw in the underlying design?

mccoyb 4y ago
I think it's young - and perhaps JAX itself is not so specialized to a specific task (but there's plenty of libraries for deep learning focused tooling, although not as mature as PyTorch). It has often been said in other threads on JAX, but it feels like a very different type of library from other AD systems -- the focus on concisely exposing/allowing users to express composable transformations seems novel! (But I may be mistaken) But in general, I would suspect youth.
time_to_smile 4y ago
I think it's better to think of JAX as a more general framework for differentiable programming and PyTorch more focused specifically on deep learning/neural networks. The beauty of JAX is that basic usage is basically a single function: `grad`. You just write whatever Python function you want and can get the derivative/gradient of it trivially. It gets a bit trickier when you need more sophisticated numeric tools like numpy/scipy, but in those cases it's just about swapping out with a JAX version of those. In this sense JAX is the spiritual success to Autograd. However the really amazing thing about JAX is that not only do you get the autodiff for basically free, you also get very good performance, and basically GPU parallelism without needing to think about it at all. PyTorch is an awesome library, but largely focus on building Neural Networks specifically. JAX should be thought of a tool that basically any Python programmer can just throw in there whenever they come across a problem that benefits from having differentiable code (which is a lot of cases once you start thinking about differentiation as a first class feature).
- jstx1 4y ago
  > I think it's better to think of JAX as a more general framework for differentiable programming and PyTorch more focused specifically on deep learning/neural networks. I don't get the point of this distinction - JAX was developed specifically for ML. What else is it being used for right now?
  - geysersam 4y ago
    My impression is that people are experimenting more with automatic differentiation for more traditional scientific computing applications. Although, I've mainly heard about Julia in that context, not Jax.
  - PartiallyTyped 4y ago
    I have used it for a numerical optimisation library that I wrote as part of a class. It’s probably also useful in population and metaheuristic scenarios where the optimisation objective can be described mathematically, allowing you to make use of GPGPUs, and if possible first and second order derivatives.

mccoyb 4y ago

I think it's young - and perhaps JAX itself is not so specialized to a specific task (but there's plenty of libraries for deep learning focused tooling, although not as mature as PyTorch). It has often been said in other threads on JAX, but it feels like a very different type of library from other AD systems -- the focus on concisely exposing/allowing users to express composable transformations seems novel! (But I may be mistaken) But in general, I would suspect youth.

time_to_smile 4y ago

I think it's better to think of JAX as a more general framework for differentiable programming and PyTorch more focused specifically on deep learning/neural networks. The beauty of JAX is that basic usage is basically a single function: `grad`. You just write whatever Python function you want and can get the derivative/gradient of it trivially. It gets a bit trickier when you need more sophisticated numeric tools like numpy/scipy, but in those cases it's just about swapping out with a JAX version of those. In this sense JAX is the spiritual success to Autograd. However the really amazing thing about JAX is that not only do you get the autodiff for basically free, you also get very good performance, and basically GPU parallelism without needing to think about it at all. PyTorch is an awesome library, but largely focus on building Neural Networks specifically. JAX should be thought of a tool that basically any Python programmer can just throw in there whenever they come across a problem that benefits from having differentiable code (which is a lot of cases once you start thinking about differentiation as a first class feature).

jstx1 4y ago
> I think it's better to think of JAX as a more general framework for differentiable programming and PyTorch more focused specifically on deep learning/neural networks. I don't get the point of this distinction - JAX was developed specifically for ML. What else is it being used for right now?
- geysersam 4y ago
  My impression is that people are experimenting more with automatic differentiation for more traditional scientific computing applications. Although, I've mainly heard about Julia in that context, not Jax.
- PartiallyTyped 4y ago
  I have used it for a numerical optimisation library that I wrote as part of a class. It’s probably also useful in population and metaheuristic scenarios where the optimisation objective can be described mathematically, allowing you to make use of GPGPUs, and if possible first and second order derivatives.