There's been some exciting work generalising transformers to data structures that aren't just pure arrays: https://glaive-research.org/2025/02/11/Generalized-Transform...
I've implemented these in Idris 2: https://github.com/bgavran/TensorType/blob/main/src/NN/Archi...