How do I get started with Jax on TPU VMs

Re:How do I get started with Jax on TPU VMs by lern_too_spel on Sunday March 31, 2024 @01:49PM (#39886332)

> When PyTorch came out it was the only option that was easy to use and debug

I'm going to go out on a limb here and suggest you didn't use PyTorch when it came out in 2016 then.

> Google had to add eager mode in TF2, ruining performance in the process

Eager mode PyTorch has worse performance than torch.compile too.

[ Parent | on HN ]

Re:How do I get started with Jax on TPU VMs by ein0p on Sunday March 31, 2024 @03:41PM (#39887363)

I used PyTorch since it came out, and at the time it outperformed TF1 on computer vision workloads quite handily, in addition to being much easier to work with. I literally watched the entire departments in one of the major research labs where I worked at the time switch to PyTorch in the span of a couple of months.

[ Parent | on HN ]