3 ms·
How does this compare with PyTorch / Tensorflow / etc.? Obviously doing heavy data processing on the GPU will have a large speedup compared to a single thread o
by xyzzy4747 4y ago
How does this compare with PyTorch / Tensorflow / etc.? Obviously doing heavy data processing on the GPU will have a large speedup compared to a single thread on the CPU.
It's almost like the author is claiming credit for creating Nvidia, when in fact he is just calling its APIs.
- m00x 4y agoI don't get that sense at all. You could do the same with tensorflow and pytorch, but in my experience, with more difficulty since they're more opinionated about how you should do your operations. JAX is definitely easier to do things that aren't on rails.
- luchris429 4y agoThe baseline we are comparing to is standard RL training that is widely used in academia. The technique mentioned in the blog post is not widely used amongst researchers. The reason we write about Jax is that doing this technique is really hard in PyTorch / Tensorflow. This is because: 1. Jax has vmap. (PyTorch does now too, but it is far more recent). 2. There are RL environments that others have written in pure Jax (see the blog post for four different repos of RL environments) 3. As m00x hints to, Jax replicates Numpy's API. This makes it way easier to use for non-neural network programming (e.g. RL environments).