For people interested in this: I've been working on an open-source libtpu implementation for https://tenstorrent.com/ accelerators as a side project, it can already run most of torchtpu, see https://github.com/pcmoritz/libtt (based on the tenstorrent software stack). It shows all the things involved in lowering torch to tile based accelerators.
Very cool! I have been using Lean to drive PJRT like so: https://github.com/brettkoonce/lean4-mlir and in theory your code could unlock another platform!