3 ms·
Google is the creator of JAX and XLA. Maybe the developer laptops have Nvidia GPUs and they do some testing there, but for Google there is literally no point in
by imtringued 1y ago
Google is the creator of JAX and XLA. Maybe the developer laptops have Nvidia GPUs and they do some testing there, but for Google there is literally no point in bothering with CUDA, pytorch or any other ecosystem strongly focused on Nvidia GPUs.
In my experience JAX is way more flexible than pytorch the moment you want to do things that aren't training ML models. E.g. you want to build an optimizer that uses the derivative of your model with respect to the input.
- kingstnap 1y agoHonestly Pytorch is weird imo, I'm surprised people love it so much. Loss.backward? Tensor.grad? Optimizer.zero grad()? With torch.no_grad()? What is with all these objects holding pointers to stuff? An ndarray is a pointer to memory and a shape my dudes. A gradient is the change in a scalar function w.r.t to some inputs.