3 ms·

> Why Should I Care About JAX? > In short - speed. For me personally, the magic of JAX is that it able to have this performance, while being as close as possi

> Why Should I Care About JAX? > In short - speed. For me personally, the magic of JAX is that it able to have this performance, while being as close as possible to having first class differentiation in Python. The latter is a far more important reason to use JAX. It can really change how you think about programming and ML. Rather than implementing a specific model, you can write up the parameterized solution to a problem then solve it. However first class differentiation ultimately isn't really useful unless you happen to also solve the speed problem. That is what makes JAX incredible. From the programming perspective JAX is to differentiable programming what Prolog is to logic programming, however Prolog has always been limited ultimately by performance problems where JAX is not.

baldeagle 5y ago
I've been really struggling to see the benefit from autodiff... I work with DAGs, where the data processes are batches that could be processed in parallel and I would love to speed that up, or to enable some gradient descent style optimization.... but I can't figure out how autodiff plays into it. Do you know of a good place to help learn how it could apply to my domain?
- SleekEagle 5y ago
  Firstly vmap() might make it easier for you from a programming standpoint. You can think of operations in terms of single instances and then autobatch them. Also, jit() is where you'll likely see your speed increases. As for autodiff - not every application will use it. What exactly are you trying to optimize? If you can parameterize your model somehow (e.g. parameterize edge weights) and then figure out some way to measure how "bad" your model is, you can use autodiff to tune your edge weights to minimize that metric. Not too familiar with DAGs but the first step is figuring out how to parameterize your model in a way that tuning the params can lead to your goal (and how to measure how close you are to that goal)
SleekEagle 5y ago
Thanks for your comment - I wanted to keep it as general as possible, and I think speed touches basically every possible use of JAX, but I definitely hear what you are saying. For DL, jvps and quick/easy Hessians (along with speed) are really the reasons to use JAX imo. I like your comment on thinking about implementing models/thinking about them as parameterized solutions

baldeagle 5y ago

I've been really struggling to see the benefit from autodiff... I work with DAGs, where the data processes are batches that could be processed in parallel and I would love to speed that up, or to enable some gradient descent style optimization.... but I can't figure out how autodiff plays into it. Do you know of a good place to help learn how it could apply to my domain?

SleekEagle 5y ago
Firstly vmap() might make it easier for you from a programming standpoint. You can think of operations in terms of single instances and then autobatch them. Also, jit() is where you'll likely see your speed increases. As for autodiff - not every application will use it. What exactly are you trying to optimize? If you can parameterize your model somehow (e.g. parameterize edge weights) and then figure out some way to measure how "bad" your model is, you can use autodiff to tune your edge weights to minimize that metric. Not too familiar with DAGs but the first step is figuring out how to parameterize your model in a way that tuning the params can lead to your goal (and how to measure how close you are to that goal)

SleekEagle 5y ago

Thanks for your comment - I wanted to keep it as general as possible, and I think speed touches basically every possible use of JAX, but I definitely hear what you are saying. For DL, jvps and quick/easy Hessians (along with speed) are really the reasons to use JAX imo. I like your comment on thinking about implementing models/thinking about them as parameterized solutions