8 ms·
To be honest, most researchers in applied ML in the bay say the opposite. If you are trying to be nimble and prototype, use pytorch. If you're trying to gain so
by lyapunova 3y ago
To be honest, most researchers in applied ML in the bay say the opposite. If you are trying to be nimble and prototype, use pytorch. If you're trying to gain some optimizations as you near deployment, rewrite in Jax.
- plumeria 3y agoWhere does Tensorflow stand in this?
- axpy906 3y agoSomewhere next to Theano, Mxnet or Caffe.
- rockinghigh 3y agoTensorflow has been falling behind since they stopped caring about backward compatibility. PyTorch is the leading framework. Jax is getting some traction at Google and was used to train Gemini.
- varelse 3y ago[dead]
- axpy906 3y agoInteresting. I’ve never heard that. I could see that argument going both ways as PyTorch has the larger ecosystem and is published the most.
- pama 3y agoInteresting perspective about possible Jax optimizations. Assuming these models are trained and deployed on non-TPU hardware, are there any real advantages in using Jax for deployment on GPU? I’d have assumed that inference is largely a solved optimization for large transformer based models (with any low hanging fruits from custom CUDA code already written) and the details are shifting towards infrastructure tradeoffs and availability of efficient GPUs. But I may be out of the loop with the latest gossip. Or do you simply mean that maybe there exist cases where TPU inference makes sense financially and using jax makes a difference?