Just wanted to share how to use Reactant+Lux on TPUs for a LLama like transformer model, it seems really nice to write even though its hard to adapt to the compact way of writting it does feel a lot better after. Still cannot take gradients though but this may be known.
Didn’t try to mesh the devices yet but it may be even better after