transformers.py contains the code to train the model. Grokking_Analysis.ipynb contains the code to load the saved checkpoints for the mainline run, calculate the progress metrics on it, and plots the figures. Non_Modular_Addition_Grokking_Tasks.ipynb contains training code for the non-modular addition experiments.