6 ms·

JAX: Numpy with Gradients, GPUs and TPUs

buildbot 8y ago
I wonder how performance compares to cupy : https://cupy.chainer.org https://cupy.chainer.org Seems a little limited in terms of supported operations for autograd compared to Chainer:https://chainer.org https://chainer.org or Flux:http://fluxml.ai http://fluxml.ai Really cool to see though! XLA/TPU support is awesome.
- alexbw 8y ago
  Actually, advanced autodiff is one of its intended points of, er, differentiation :). The authors wrote the original Autograd package [0], released in 2014, that led to “autograd” becoming used as a generic term in PyTorch and other packages. JAX has all of the autodiff operations that Autograd does, including `grad`, `vjp`, `jvp`, etc. We’re working on the number of supported NumPy ops, which is limited right now, but it’s early days. Try it out, we’re really excited to see what you build with it! 0: https://github.com/hips/autograd https://github.com/hips/autograd
chrisaycock 8y ago
For anyone else outside of machine learning who was wondering what all of this is, here is my best explanation: The inferencing phase of a neural network attempts to minimize error or loss as defined by the user. This is done by iteratively applying gradient descent to the error function. Thus, the error function must have a known derivative, which can be difficult if the function has loops and conditionals. Autograd is a software package that produces derivatives of a function automatically. It creates a computation graph of the user-defined code from which it can determine the derivative: https://en.wikipedia.org/wiki/Automatic_differentiation https://en.wikipedia.org/wiki/Automatic_differentiation XLA is a JIT from TensorFlow that compiles common array functions. The JAX project from this GitHub page brings JIT optimizations to Autograd's automatic differentiation. That will speed-up the error function when inferring a neural network. Neat! (I would be grateful for any corrections to my above explanation as I am not an expert in ML.)
- dgacmu 8y ago
  Small correction: XLA is a compiler in the more general sense, not a JIT specifically. (It's a domain-specific compiler focused on linear algebra.). JAX uses XLA to do JIT compilation.
- twtw 8y ago
  Another small correction: every instance of "inference" in your comment should probably be replaced with "training." It's the training phase the involves running gradient descent of various flavors to optimize the network parameters.
  - wenc 8y ago
    This is an important point. Inferencing means “to predict” (I’m not sure when this terminology became popular; a few years ago most of us were just using the word predict) Once trained, a model no longer requires derivatives. It’s more or less a function evaluation, which can be done on plain CPUs.
    - conistonwater 8y ago
      It's from statistical inference, e.g., when the goal is to find the values of a model's parameters that match the sample. So if the model is y = f(x, params), inference gives you params, and prediction gives you y for a value of x that you haven't seen before. Also, shouldn't the verb be inferring?
      - wenc 8y ago
        Except in AI terminology[1], it seems that inference doesn't mean outputting params. It means outputting y. Yes, I believe the verb is "inferring". [1] https://blogs.nvidia.com/blog/2016/08/22/difference-deep-learning-training-inference-ai/ https://blogs.nvidia.com/blog/2016/08/22/difference-deep-lea...
        conistonwater 8y ago
hyperpallium 8y ago
> JAX can automatically differentiate native Python ... functions sympy can differentiate functions, but they have to be set up properly. How can JAX differentiate native functions? (Or do they mean numerical differentiation, like a finite difference estimation?)
- currymj 8y ago
  Neither, they mean automatic differentiation [1], which is not symbolic like SymPy or Mathematica, but also is not finite differences. Often, the interface for doing this is to just extend an existing numerical or array type and overload all the arithmetic operations to keep track of the gradients. Then ordinary code for numerical computations will "just work". That's basically what they've done here, except with a sophisticated compiler. 1: https://en.wikipedia.org/wiki/Automatic_differentiation https://en.wikipedia.org/wiki/Automatic_differentiation
  - RivieraKid 8y ago
    How's AD different from simply calling the function with a slightly perturbed input and observing how much the output changes? I assume it's more efficient with multivariable functions because the aforementioned method requires one call per parameter?
    - shoyer 8y ago
      It is indeed much more efficient. In general, you can evaluate a scalar function and its gradient with less than twice the effort required to compute the scalar function alone -- regardless of the number of parameters.
    - sampo 8y ago
      > How's AD different from simply calling the function with a slightly perturbed input and observing how much the output changes? AD gives a more precise result, as AD calculates the value of the derivative at the point you want without any perturbed input. AD is also faster, as calculating the value of the derivative usually requires about the same number of elementary operations as calculating the value of the function. Whereas with finite difference you need to calculate the function twice.
snackematician 8y ago
Very cool! I love autograd, it had tape-based autodiff way before pytorch, and the way it wraps numpy is much more convenient than tensorflow/pytorch. Been wanting GPU support in autograd for years now, so am very happy to see this. I have some academic software (https://github.com/popgenmethods/momi2 https://github.com/popgenmethods/momi2) that uses autograd, was planning to port it to pytorch since it's better supported/maintained, but now I'll have to consider jax. Though I'm a little worried about the maturity of the project, seems like the numpy/scipy coverage is not all the way there yet. Then again, it would be fun to contribute back to JAX, I did contribute a couple PRs to autograd back in the day so I think I could jump right into it...
duckworthd 8y ago
My $0.02. I've been using JAX for a while now. A paper I'm an author on (https://arxiv.org/abs/1806.09597 https://arxiv.org/abs/1806.09597, a follow-up to https://news.ycombinator.com/item?id=18633215 https://news.ycombinator.com/item?id=18633215) resulted in an algorithm that required taking second-derivatives on a per-example basis. This is extremely difficult in TensorFlow, but with JAX it was a 2-liner. Even better, it's _super_ fast, thanks to XLA's compile-to-GPU and JAX's auto-batching mechanics. I highly recommend JAX to power users. It's nowhere near as feature-complete from a neural network sense as, say, PyTorch, but it is very good at what it does, and its core developers are second to none in responsiveness.
- albertzeyer 8y ago
  Why is that difficult in TensorFlow? Wouldn’t you just call tf.gradients on the output of tf.gradients? At least that is what I usually do and it’s very simple.
  - duckworthd 8y ago
    There are two reasons, 1) Take a closer look at tf.hessian()'s implementation. Notice the tf.while() loop, calculating tf.gradient() for each coordinate of the first gradient? That doesn't happen in parallel! Plus, you need to know about TensorArrays! 2) Per-example gradients means an outside tf.while() loop over each example. Another linear slow down! The difficulty I refer to is in wrangling conditionals in TF, and the trickiness of obtaining partial derivatives (rather than total derivatives) which JAX makes trivial. There may be a
    - albertzeyer 8y ago
      tf.while code will run in parallel as far as possible (if it does not depend on each other) (up to some configurable degree).
      - duckworthd 8y ago
        Yes! Though this answer has some subtlety. tf.while() will run several iterations in parallel but this is not the same as _batching_ those same iterations. For that, you'll need to use the experimental parallel_for feature [0]. Using this should get you into roughly the same speed as JAX. Tricky! [0]: https://github.com/tensorflow/tensorflow/tree/b3e007394680801113f492fa1f5a9784e8502f19/tensorflow/python/ops/parallel_for https://github.com/tensorflow/tensorflow/tree/b3e00739468080...
p1esk 8y ago
How does it compare to CuPy (in terms of Numpy compatibility)?
lostmsu 8y ago
What does it do, if the function is not differentiable? Most of the commonly seen functions in real code aren't.
whoisnnamdi 8y ago
This is very cool, and I can see all sorts of use cases where a tool like this could be valuable. Definitely will try to keep in mind that a tool for fast differentiation of arbitrary functions exists out in the world when starting my next project Thanks for posting!

I wonder how performance compares to cupy : https://cupy.chainer.org https://cupy.chainer.org Seems a little limited in terms of supported operations for autograd compared to Chainer:https://chainer.org https://chainer.org or Flux:http://fluxml.ai http://fluxml.ai Really cool to see though! XLA/TPU support is awesome.

alexbw 8y ago
Actually, advanced autodiff is one of its intended points of, er, differentiation :). The authors wrote the original Autograd package [0], released in 2014, that led to “autograd” becoming used as a generic term in PyTorch and other packages. JAX has all of the autodiff operations that Autograd does, including `grad`, `vjp`, `jvp`, etc. We’re working on the number of supported NumPy ops, which is limited right now, but it’s early days. Try it out, we’re really excited to see what you build with it! 0: https://github.com/hips/autograd https://github.com/hips/autograd

chrisaycock 8y ago

For anyone else outside of machine learning who was wondering what all of this is, here is my best explanation: The inferencing phase of a neural network attempts to minimize error or loss as defined by the user. This is done by iteratively applying gradient descent to the error function. Thus, the error function must have a known derivative, which can be difficult if the function has loops and conditionals. Autograd is a software package that produces derivatives of a function automatically. It creates a computation graph of the user-defined code from which it can determine the derivative: https://en.wikipedia.org/wiki/Automatic_differentiation https://en.wikipedia.org/wiki/Automatic_differentiation XLA is a JIT from TensorFlow that compiles common array functions. The JAX project from this GitHub page brings JIT optimizations to Autograd's automatic differentiation. That will speed-up the error function when inferring a neural network. Neat! (I would be grateful for any corrections to my above explanation as I am not an expert in ML.)

dgacmu 8y ago
Small correction: XLA is a compiler in the more general sense, not a JIT specifically. (It's a domain-specific compiler focused on linear algebra.). JAX uses XLA to do JIT compilation.
twtw 8y ago
Another small correction: every instance of "inference" in your comment should probably be replaced with "training." It's the training phase the involves running gradient descent of various flavors to optimize the network parameters.
- wenc 8y ago
  This is an important point. Inferencing means “to predict” (I’m not sure when this terminology became popular; a few years ago most of us were just using the word predict) Once trained, a model no longer requires derivatives. It’s more or less a function evaluation, which can be done on plain CPUs.
  - conistonwater 8y ago
    It's from statistical inference, e.g., when the goal is to find the values of a model's parameters that match the sample. So if the model is y = f(x, params), inference gives you params, and prediction gives you y for a value of x that you haven't seen before. Also, shouldn't the verb be inferring?
    - wenc 8y ago
      Except in AI terminology[1], it seems that inference doesn't mean outputting params. It means outputting y. Yes, I believe the verb is "inferring". [1] https://blogs.nvidia.com/blog/2016/08/22/difference-deep-learning-training-inference-ai/ https://blogs.nvidia.com/blog/2016/08/22/difference-deep-lea...
      - conistonwater 8y ago

hyperpallium 8y ago

> JAX can automatically differentiate native Python ... functions sympy can differentiate functions, but they have to be set up properly. How can JAX differentiate native functions? (Or do they mean numerical differentiation, like a finite difference estimation?)

currymj 8y ago
Neither, they mean automatic differentiation [1], which is not symbolic like SymPy or Mathematica, but also is not finite differences. Often, the interface for doing this is to just extend an existing numerical or array type and overload all the arithmetic operations to keep track of the gradients. Then ordinary code for numerical computations will "just work". That's basically what they've done here, except with a sophisticated compiler. 1: https://en.wikipedia.org/wiki/Automatic_differentiation https://en.wikipedia.org/wiki/Automatic_differentiation
- RivieraKid 8y ago
  How's AD different from simply calling the function with a slightly perturbed input and observing how much the output changes? I assume it's more efficient with multivariable functions because the aforementioned method requires one call per parameter?
  - shoyer 8y ago
    It is indeed much more efficient. In general, you can evaluate a scalar function and its gradient with less than twice the effort required to compute the scalar function alone -- regardless of the number of parameters.
  - sampo 8y ago
    > How's AD different from simply calling the function with a slightly perturbed input and observing how much the output changes? AD gives a more precise result, as AD calculates the value of the derivative at the point you want without any perturbed input. AD is also faster, as calculating the value of the derivative usually requires about the same number of elementary operations as calculating the value of the function. Whereas with finite difference you need to calculate the function twice.

snackematician 8y ago

Very cool! I love autograd, it had tape-based autodiff way before pytorch, and the way it wraps numpy is much more convenient than tensorflow/pytorch. Been wanting GPU support in autograd for years now, so am very happy to see this. I have some academic software (https://github.com/popgenmethods/momi2 https://github.com/popgenmethods/momi2) that uses autograd, was planning to port it to pytorch since it's better supported/maintained, but now I'll have to consider jax. Though I'm a little worried about the maturity of the project, seems like the numpy/scipy coverage is not all the way there yet. Then again, it would be fun to contribute back to JAX, I did contribute a couple PRs to autograd back in the day so I think I could jump right into it...

duckworthd 8y ago

My $0.02. I've been using JAX for a while now. A paper I'm an author on (https://arxiv.org/abs/1806.09597 https://arxiv.org/abs/1806.09597, a follow-up to https://news.ycombinator.com/item?id=18633215 https://news.ycombinator.com/item?id=18633215) resulted in an algorithm that required taking second-derivatives on a per-example basis. This is extremely difficult in TensorFlow, but with JAX it was a 2-liner. Even better, it's _super_ fast, thanks to XLA's compile-to-GPU and JAX's auto-batching mechanics. I highly recommend JAX to power users. It's nowhere near as feature-complete from a neural network sense as, say, PyTorch, but it is very good at what it does, and its core developers are second to none in responsiveness.

albertzeyer 8y ago
Why is that difficult in TensorFlow? Wouldn’t you just call tf.gradients on the output of tf.gradients? At least that is what I usually do and it’s very simple.
- duckworthd 8y ago
  There are two reasons, 1) Take a closer look at tf.hessian()'s implementation. Notice the tf.while() loop, calculating tf.gradient() for each coordinate of the first gradient? That doesn't happen in parallel! Plus, you need to know about TensorArrays! 2) Per-example gradients means an outside tf.while() loop over each example. Another linear slow down! The difficulty I refer to is in wrangling conditionals in TF, and the trickiness of obtaining partial derivatives (rather than total derivatives) which JAX makes trivial. There may be a
  - albertzeyer 8y ago
    tf.while code will run in parallel as far as possible (if it does not depend on each other) (up to some configurable degree).
    - duckworthd 8y ago
      Yes! Though this answer has some subtlety. tf.while() will run several iterations in parallel but this is not the same as _batching_ those same iterations. For that, you'll need to use the experimental parallel_for feature [0]. Using this should get you into roughly the same speed as JAX. Tricky! [0]: https://github.com/tensorflow/tensorflow/tree/b3e007394680801113f492fa1f5a9784e8502f19/tensorflow/python/ops/parallel_for https://github.com/tensorflow/tensorflow/tree/b3e00739468080...

p1esk 8y ago

How does it compare to CuPy (in terms of Numpy compatibility)?

lostmsu 8y ago

What does it do, if the function is not differentiable? Most of the commonly seen functions in real code aren't.

whoisnnamdi 8y ago

This is very cool, and I can see all sorts of use cases where a tool like this could be valuable. Definitely will try to keep in mind that a tool for fast differentiation of arbitrary functions exists out in the world when starting my next project Thanks for posting!