4 ms·
Is there any good way to program TPUs for non-ML work? The only official way appears to involve encoding the compute operation into a TF graph and loading that.
by sxp 3y ago
Is there any good way to program TPUs for non-ML work? The only official way appears to involve encoding the compute operation into a TF graph and loading that. Is there a better way to directly use the hardware?
- counters 3y agoJust write your code in JAX and run as usual?
- pjmlp 3y agoLooking forward to see Nanite implementation in JAX.
- lordswork 3y agoAFAIK, XLA is the only compiler that can target TPUs, so anything written for TPUs would need to lower to HLO. That said, why would you want to to run non-ML workloads on such an ML-specialized chip?
- sxp 3y agoI was interested in seeing if raytracing or similar embarrassingly parallel workloads could be ported to a TPU.