Code for running the transformers in the ICML 2021 paper "Thinking Like Transformers" - View it on GitHub