Code for paper : How to think step-by-step: A mechanistic understanding of chain-of-thought reasoning