4 ms·

I haven't worked with TPUs, but my understanding is that they are pretty plug-n-play for Google frameworks (JAX, TF) but is also pretty simple to use with PyTor

• shakna 1y ago
  If only things were simple [0]. [0] https://github.com/pytorch/xla/issues/9189 https://github.com/pytorch/xla/issues/9189
  • jszymborski 1y ago
    Certainly not plug-and-play, and perhaps this is just one of many bugs, but if it isn't I can't imagine a scenario where LazyLinear layers are essential.
• que-encrypt 1y ago
  pytorch xla is barely supported in the pytorch ecosystem (for instance, pytorch lightning still doesn't easily support tpu pods, with only a singular short page about google colab v2-8 tpus that is out of date. Then there are the various libraries/implementations with pytorch that have a .cuda(), etc. More limitations at: https://lightning.ai/docs/pytorch/stable/accelerators/tpu_faq.html https://lightning.ai/docs/pytorch/stable/accelerators/tpu_fa...). I haven't worked with tensorflow, but I've heard it's a pain even when using gpus. JAX is the real deal, and does make my code transferrable between GPUs/TPUs relatively easily (excluding any custom pallas kernels for flash vs splash attention, but this is usually not a massive code change). However, with JAX, there are often not a bunch of pre-existing implementations due to network effects, etc.