SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Theorem 3 gives a causal map that can be computed entirely using Monarch matrices - enforcing causality with sub-quadratic scaling. The main technical ingredient in proving the above result is that the product  $q_{j}(Z)q_{j^{\prime}}(Z)$  can be written as a linear combination of  $q_{a}(Z)$  for  $j + j^{\prime} \leq a < N$  (this uses the above specified properties on the minimum and maximum degrees of  $q_{j}(Z)$ ). This in turn implies that the term  $k_{j^{\prime}}u_{j}q_{j}(Z)q_{j^{\prime}}(Z)$  only contributes to the coefficients of "higher order" basis polynomials  $q_{a}(Z)$  for  $a \geq j + j^{\prime}$  in the product  $k(Z)u(Z)$ , which is needed for causality. Figure 2 gives an example of restricted polynomials generating a causal map.
We compare MONARCH MIXER to Transformers on three tasks where Transformers have been dominant: BERT-style non-causal masked language modeling, ViT-style image classification, and GPT-style causal language modeling. In each, we show that we can match Transformers in quality using neither attention nor MLPs. We additionally evaluate wall-clock speedups against strong Transformer baselines in the BERT setting. Additional experiments on speech and alternative architectures are given in Appendix B, and experimental details are given in Appendix C.
5.1 Non-Causal Language Modeling
We introduce M2-BERT, an M2-based architecture for non-causal language modeling. M2-BERT acts as a drop-in replacement for BERT-style language models [21], which are a workhorse application of the Transformer architecture [1, 39, 40, 45, 48, 49, 52, 56, 85, 89]. We train M2-BERT using masked language modeling over C4 [66] with the bert-base-uncased tokenizer.
Figure 3: M2-BERT uses Monarch matrices to create a bidirectional gated long convolution in the sequence mixer, and uses Monarch matrices to replace the linear layers in the dimension mixer.
M2-BERT starts with a Transformer backbone and replaces the attention and MLPs with M2 layers, shown in Figure 3. In the sequence mixer, we replace attention with bidirectional gated convolutions with a residual convolution (Figure 3 left). To recover convolutions, we set the Monarch matrices to DFT and inverse DFT matrices. Following [27, 64], we also add short depthwise convolutions after the projections. In the dimension mixer, we replace the two dense matrices in MLPs with learned block-diagonal matrices (Monarch matrix of order  $1$ ,  $b = 4$ ). We pretrain two M2-BERT-base models, at 80M and 110M, and two M2-BERT-large models, at 260M and 341M. These are equivalent to BERT-base and BERT-large, respectively.
Downstream GLUE Scores First, we evaluate M2-BERT models on downstream fine-tuning compared to BERT-base and BERT-large from [20]. We take the pretrained models and fine-tune them on BERT, following the procedure in [38]. Table 3 shows performance for BERT-base equivalent models, and Table 4 shows performance for BERT-large equivalent models. M2-BERT-base can
Table 3: Average GLUE Score for M2-BERT-base compared to BERT-base [20], along with change in parameters and GLUE score.
Model	GLUE Score	Δ Params	Δ GLUE Score
BERT-base (110M)	79.6	-0%	+0.0
M2-BERT-base (80M)	79.9	-27%	+0.3
M2-BERT-base (110M)	80.9	-0%	+1.3
Table 4: Average GLUE Score for M2-BERT-large compared to BERT-large [20], along with change in parameters and GLUE score.
Model	GLUE Score	Δ Params	Δ GLUE Score
BERT-large (340M)	82.1	-0%	+0.0
M2-BERT-large (260M)	82.2	-24%	+0.1
M2-BERT-large (341M)	82.8	+0.2%	+0.7
match BERT-base in GLUE quality with  $27\%$  fewer parameters—or outperform BERT-base in quality by 1.3 points when parameter matched. M2-BERT-large matches BERT-large with  $24\%$  fewer parameters, and outperforms by 0.7 points when parameter matched.
GPU Throughput by Sequence Length Next, we evaluate throughput of M2-BERT models by sequence length, compared to HuggingFace implementations of BERT, as well as optimized implementations of BERT running FlashAttention [15]. Table 5 shows forward throughput for BERT-base equivalent models, and the appendix shows throughput for BERT-large (where the performance trends are similar). Inference times are reported in tokens/ms on an A100-40GB GPU. M2-BERT-base achieves higher throughput than even highly-optimized BERT models, and up to  $9.1 \times$  faster throughput than a standard HuggingFace implementation at sequence length 4K.
CPU Inference Latency Finally, we report CPU inference latency for M2-BERT-base (80M) compared to BERT-base, running direct PyTorch implementations for both. In short sequences, the impacts of data locality still dominate the FLOP reduction, and operations such as filter generation (which are not present in BERT) pay a higher cost. Starting at sequences 1K and longer, M2-BERT-base starts to have speedup over BERT-base, up to  $6.5 \times$  at sequence length 8K. We believe further optimization and applying IO-aware principles can further improve CPU performance.
5.2 Image Classification
To validate that our methods generalize to images as well as language for non-causal modeling, we next evaluate M2 on image classification. We compare M2 to ViT-style models and recent work, HyenaViT-b [64], which uses gated long convolutions to replace the attention layers in ViT-b. In our work, M2-ViT builds off HyenaViT-b and replaces the long convolutions with the M2 operator in Equation 2 (again setting the Monarch matrices to the DFT and inverse DFT). We replace the MLP blocks in HyenaViT-b with block-diagonal matrices, similarly to M2-BERT. Appendix B additionally compares M2 to the Swin-family of architectures [50, 51].
Table 7 shows the performance of MONARCH MIXER against ViT-b, HyenaViT-b, and ViT-b-Monarch (which replaces the MLP blocks of standard ViT-b with Monarch matrices) on ImageNet-1k. MONARCH MIXER outperforms the other models with only half the parameters of the original ViT-s model. Surprisingly, MONARCH MIXER also outperforms ResNet-152, with fewer parameters—even though the latter was explicitly designed for ImageNet performance.
5.3 Causal Language Modeling
GPT-style causal language modeling is a critical application for Transformers [6, 31, 43]. We introduce M2-GPT, a M2-based architecture for causal language modeling. For the sequence mixer, M2-GPT combines the convolutional filter from Hyena [64], the state-of-the-art attention-free
Table 5: Throughput in tokens/ms by context length for M2-BERT-base (80M) compared to BERT-base.
Model	512	1024	2048	4096	8192
HF BERT-base (110M)	206.1	130.8	71.3	39.0	OOM
FlashAttention BERT-base (110M)	367.4	350.1	257.2	179.1	102.4
M2-BERT-base (80M)	386.3	380.7	378.9	353.9	320.1
M2 Speedup over HF BERT-base (110M)	1.9×	2.9×	5.2×	9.1×	-
Table 6: CPU inference latency in milliseconds with a batch size of 1 at varied input sequence lengths. Measurements averaged over 10 examples on a 48 vCPU, 96 GB RAM instance from the GCP n2-standard-48 series, which runs Intel Cascade Lake processors. This is based on the protocol in [29].
Model	512	1024	2048	4096	8192
BERT-base (110M)	182	389	918	2660	11820
M2-BERT-base (80M)	289	361	651	948	1820
Speedup	0.6×	1.1×	1.4×	2.8×	6.5×
language model, with parameter sharing across multiple heads from H3 [27]. We use the causal parameterization of Equation 2 to replace the FFT in these architectures, and we remove the MLP layers entirely. The resulting architecture is entirely attention- and MLP-free.
We pretrain M2-GPT on the PILE, a standard dataset for causal language modeling. Following prior work [28, 64], we train models at two model sizes, with varying amounts of training data—decaying the learning rate appropriately for each experiment. Table 8 shows the results. Even though our model is attention- and MLP-free, it outperforms both Transformers and Hyena in perplexity on pretraining. These results suggest that radically different architectures than Transformers may be performant on causal language modeling.
**Long Convolutions** Recent work proposes to use long convolution layers as a replacement for the Transformer attention layers in sequence modeling [28, 64, 67-69]. Many of these models rely on the FFT convolution theorem to compute the long convolutions. We build on the insights in many of these architectures in constructing our M2 architectures, and additionally replaces the FFT operations with Monarch matrices.
Our work is also related to a rich literature in convolutions in other bases, such as Chebyshev bases [82] or orthogonal polynomial bases [34]. These approaches have analogues in our multivariate analysis; replacing the basis polynomials of the Monarch matrices in MONARCH MIXER may be able to approximate some of these operations. An interesting question for future work would be to study how well our techniques and concerns about causality and hardware utilization translate to these alternative convolution bases.
Optimization of deep learning primitives There is a rich history of the optimization of deep learning primitives, as accelerating their performance can yield substantial savings in compute and cost for large models. There are many approaches to speed up these operations, but they usually either reduce data movement or compute.
Reducing data movement: In many applications, the major bottleneck is the storage and movement of large amounts of memory. One popular approach to reducing data movement is checkpointing, wherein one stores fewer intermediate results and recomputes the others on-the-fly where they are needed, trading additional compute for memory [46, 78]. Another approach is kernel fusion, wherein algorithms initially described as sequential steps can often be fused in ways that improve their properties. For example, it is generally faster to implement a dot-product through a multiply-
Table 7: Accuracy on ImageNet-1k. ResNet-152 provided for reference.
Model	Top-1%	Top-5%	Description
ResNet-152 (60M)	78.6	94.3	ConvNet, MLP
ViT-b (87M)	78.5	93.6	Attention, MLP
ViT-b + Monarch (33M)	78.9	94.2	Attention, MLP-Free
HyenaViT-b (88M)	78.5	93.6	Attention-Free, MLP
M2-ViT-b (45M)	79.5	94.5	Attention-Free, MLP-Free
Table 8: Perplexity on the PILE when trained for different numbers of tokens.
Model	5B	10B	15B	Description
Transformer (125M)	13.3	11.9	11.2	Attention, MLP
Hyena (155M)	13.1	11.8	11.1	Attention-Free, MLP
M2-GPT (145M)	12.9	11.6	10.9	Attention-Free, MLP-Free
Transformer (355M)	11.4	9.8	9.1	Attention, MLP
Hyena (360M)	11.3	9.8	9.2	Attention-Free, MLP
M2-GPT (360M)	11.0	9.6	9.0	Attention-Free, MLP-Free
accumulate rather than first multiplying and then accumulating. Recently, libraries such as PyTorch 2.0 [63] have added kernel fusion capabilities, although the very best performance usually still arises from handwritten kernels. Third, in order to better exploit memory locality, it is often fastest to load small blocks of memory, do intensive computation on them, and then write the results a tile at a time [83]. Finally, many algorithms also have hand-optimizations that can remove unnecessary computation or memory accesses [55].
Efficient algorithms usually make use of a combination of these techniques. For example, FlashAttention [15] uses all four to dramatically decrease both the latency and memory consumption of multi-head attention. Though we have made a modest effort to implement MONARCH MIXER efficiently, we think it likely that MONARCH MIXER could be further optimized by these techniques.
Reducing flops: A first target for optimization is the multi-layer perceptron (MLP), owing to its ubiquity. A variety of structured sparse factorizations exist, many of which we draw on in this work [7, 11, 14, 16, 17, 19, 26, 91]. Attention is also a popular target for optimization. Recently, a plethora of sub-quadratic approximations of attention have emerged, that aim to approximate attention to reduce its quadratic complexity. Some methods rely on sparsification, relying on the fact that the attention matrix is extremely sparse at long sequence lengths [2, 23, 24, 42, 53]. Others use low-rank approximations of the attention matrix [13, 79, 91] or kernel methods instead [9, 41]. A subset uses a combination of these techniques, such as [8, 72]. Finally, a third category of methods [27, 64] aims to replace attention entirely, relying on state-space models [33].
We explore MONARCH MIXER (M2), a new architecture that is sub-quadratic in both sequence length and model dimension and is hardware-efficient on modern accelerators. We motivate M2 from both theoretical and systems performance perspectives and conduct a preliminary proof-of-concept investigation into performance on masked language modeling, image classification, and causal language modeling.
While our initial results are promising, our work is only a first step in this direction. The M2 layer can likely be further optimized with systems optimization techniques such as kernel fusion. Our work has also not been optimized for inference like more well-established models such as Transformers, or even more recent models such as state space models. It also remains to be seen whether M2 layers can have as widespread applicability as Transformers. We hope that these can be fruitful directions for future work.
A discussion of broader impacts can be found in the Appendix.
We gratefully acknowledge the support of DARPA under Nos. FA86501827865 (SDH) and FA86501827882 (ASED); NIH under No. U54EB020405 (Mobilize), NSF under Nos. CCF1763315 (Beyond Sparsity), CCF1563078 (Volume to Velocity), and 1937301 (RTML); ONR under No. N000141712266 (Unifying Weak Supervision); the Moore Foundation, NXP, Xilinx, LETI-CEA, Intel, IBM, Microsoft, NEC, Toshiba, TSMC, ARM, Hitachi, BASF, Accenture, Ericsson, Qualcomm, Analog Devices, the Okawa Foundation, American Family Insurance, Google Cloud, Microsoft Azure, Swiss Re, Brown Institute for Media Innovation, Department of Defense (DoD) through the National Defense Science and Engineering Graduate Fellowship (NDSEG) Program, Fannie and John Hertz Foundation, National Science Foundation Graduate Research Fellowship Program, Texas Instruments Stanford Graduate Fellowship in Science and Engineering, and members of the Stanford DAWN project: Teradata, Facebook, Google, Ant Financial, NEC, VMWare, andInfosys. The U.S. Government is authorized to reproduce and distribute reprints for Governmental purposes notwithstanding any copyright notation thereon. Any opinions, findings, and conclusions or recommendations expressed in this material are those of the authors and do not necessarily reflect the views, policies, or endorsements, either expressed or implied, of DARPA, NIH, ONR, or the U.S. Government. JG and AR's work is supported by NSF grant# CCF-2247014. IJ's work is supported by an NSF Graduate Fellowship.
[1] Iz Beltagy, Kyle Lo, and Arman Cohan. Scibert: A pretrained language model for scientific text. arXiv preprint arXiv:1903.10676, 2019.
[2] Iz Beltagy, Matthew E Peters, and Arman Cohan. Longformer: The long-document transformer. arXiv preprint arXiv:2004.05150, 2020.
[3] Emily M. Bender, Timnit Gebru, Angelina McMillan-Major, and Shmargaret Shmitchell. On the dangers of stochastic parrots: Can language models be too big? In Proceedings of the 2021 ACM Conference on Fairness, Accountability, and Transparency, FAccT '21, page 610-623, New York, NY, USA, 2021. Association for Computing Machinery.
[4] Lucas Beyer, Xiaohua Zhai, and Alexander Kolesnikov. Better plain vit baselines for imagenet-1k. arXiv preprint arXiv:2205.01580, 2022.
[5] Rishi Bommasani, Drew A Hudson, Ehsan Adeli, Russ Altman, Simran Arora, Sydney von Arx, Michael S Bernstein, Jeannette Bohg, Antoine Bosselut, Emma Brunskill, et al. On the opportunities and risks of foundation models. arXiv preprint arXiv:2108.07258, 2021.
[6] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020.
[7] Beidi Chen, Tri Dao, Kaizhao Liang, Jiaming Yang, Zhao Song, Atri Rudra, and Christopher Ré. Pixelated butterfly: Simple and efficient sparse training for neural network models. 2021.
[8] Beidi Chen, Tri Dao, Eric Winsor, Zhao Song, Atri Rudra, and Christopher Ré. Scatterbrain: Unifying sparse and low-rank attention. In Advances in Neural Information Processing Systems (NeurIPS), 2021.
[9] Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Davis, Afroz Mohiuddin, Lukasz Kaiser, et al. Rethinking attention with performers. arXiv preprint arXiv:2009.14794, 2020.
[10] Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, et al. Palm: Scaling language modeling with pathways. arXiv preprint arXiv:2204.02311, 2022.
[11] James W Cooley and John W Tukey. An algorithm for the machine calculation of complex fourier series. Mathematics of computation, 19(90):297-301, 1965.
[46] Mitsuru Kusumoto, Takuya Inoue, Gentaro Watanabe, Takuya Akiba, and Masanori Koyama. A graph theoretic framework of recomputation algorithms for memory-efficient backpropagation. Advances in Neural Information Processing Systems, 32, 2019.
[47] Lagrange polynomial. Lagrange polynomial — Wikipedia, the free encyclopedia, 2005. https://en.wikipedia.org/wiki/Lagrange_polynomial.
[48] Jinhyuk Lee, Wonjin Yoon, Sungdong Kim, Donghyeon Kim, Sunkyu Kim, Chan Ho So, and Jaewoo Kang. Biobert: a pre-trained biomedical language representation model for biomedical text mining. Bioinformatics, 36(4):1234-1240, 2020.
[49] Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. Roberta: A robustly optimized bert pretraining approach. arXiv preprint arXiv:1907.11692, 2019.
[50] Ze Liu, Han Hu, Yutong Lin, Zhuliang Yao, Zhenda Xie, Yixuan Wei, Jia Ning, Yue Cao, Zheng Zhang, Li Dong, Furu Wei, and Baining Guo. Swin transformer v2: Scaling up capacity and resolution. In International Conference on Computer Vision and Pattern Recognition (CVPR), 2022.
[51] Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo. Swin transformer: Hierarchical vision transformer using shifted windows. In Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV), 2021.
[52] Xiaofei Ma, Zhiguo Wang, Patrick Ng, Ramesh Nallapati, and Bing Xiang. Universal text representation from bert: An empirical study. arXiv preprint arXiv:1910.07973, 2019.
[53] Xuezhe Ma, Xiang Kong, Sinong Wang, Chunting Zhou, Jonathan May, Hao Ma, and Luke Zettlemoyer. Luna: Linear unified nested attention. Advances in Neural Information Processing Systems, 34:2441-2453, 2021.
[54] Xuezhe Ma, Chunting Zhou, Xiang Kong, Junxian He, Liangke Gui, Graham Neubig, Jonathan May, and Luke Zettlemoyer. Mega: moving average equipped gated attention. arXiv preprint arXiv:2209.10655, 2022.
[55] Maxim Milakov and Natalia Gimelshein. Online normalizer calculation for softmax. arXiv preprint arXiv:1805.02867, 2018.
[56] Derek Miller. Leveraging bert for extractive text summarization on lectures. arXiv preprint arXiv:1906.04165, 2019.
[57] Marcin Moczulski, Misha Denil, Jeremy Appleyard, and Nando de Freitas. Acdc: A structured efficient linear layer. arXiv preprint arXiv:1511.05946, 2015.
[58] NVIDIA. Nvidia Tesla V100 GPU architecture, 2017.
[59] NVIDIA. Nvidia A100 tensor core GPU architecture, 2020.
[60] NVIDIA. Nvidia H100 tensor core GPU architecture, 2022.
[61] NVIDIA. cuBLAS, 2023.
[62] OpenAI. Gpt-4 technical report, 2023.
[63] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high-performance deep learning library. Advances in neural information processing systems, 32, 2019.
[64] Michael Poli, Stefano Massaroli, Eric Nguyen, Daniel Y Fu, Tri Dao, Stephen Baccus, Yoshua Bengio, Stefano Ermon, and Christopher Ré. Hyena hierarchy: Towards larger convolutional language models. International Conference on Machine Learning, 2023.
[82] Kuan Xu and Ana F. Loureiro. Spectral approximation of convolution operators. SIAM Journal on Scientific Computing, 40(4):A2336-A2355, 2018.
[83] Yufan Xu, Saurabh Raje, Atanas Rountev, Gerald Sabin, Aravind Sukumaran-Rajam, and P Sadayappan. Training of deep learning pipelines on memory-constrained gpus via segmented fused-tiled execution. In Proceedings of the 31st ACM SIGPLAN International Conference on Compiler Construction, pages 104–116, 2022.
[84] Lili Yu, Daniel Simig, Colin Flaherty, Armen Aghajanyan, Luke Zettlemoyer, and Mike Lewis. Megabyte: Predicting million-byte sequences with multiscale transformers, 2023.
[85] Shanshan Yu, Jindian Su, and Da Luo. Improving bert-based text classification with auxiliary sentence and domain knowledge. IEEE Access, 7:176600-176612, 2019.
[86] Li Yuan, Yunpeng Chen, Tao Wang, Weihao Yu, Yujun Shi, Zi-Hang Jiang, Francis EH Tay, Jiashi Feng, and Shuicheng Yan. Tokens-to-token vit: Training vision transformers from scratch onImagenet. In Proceedings of the IEEE/CVF international conference on computer vision, pages 558-567, 2021.
[87] Sangdoo Yun, Dongyoon Han, Seong Joon Oh, Sanghyuk Chun, Junsuk Choe, and Youngjoon Yoo. Cutmix: Regularization strategy to train strong classifiers with localizable features. In Proceedings of the IEEE/CVF international conference on computer vision, pages 6023-6032, 2019.
[88] Hongyi Zhang, Moustapha Cisse, Yann N Dauphin, and David Lopez-Paz. mixup: Beyond empirical risk minimization. arXiv preprint arXiv:1710.09412, 2017.
[89] Tianyi Zhang, Varsha Kishore, Felix Wu, Kilian Q Weinberger, and Yoav Artzi. Bertscore: Evaluating text generation with bert. arXiv preprint arXiv:1904.09675, 2019.
[90] Zhun Zhong, Liang Zheng, Guoliang Kang, Shaozi Li, and Yi Yang. Random erasing data augmentation. In Proceedings of the AAAI conference on artificial intelligence, volume 34, pages 13001-13008, 2020.
[91] Chen Zhu, Wei Ping, Chaowei Xiao, Mohammad Shoeybi, Tom Goldstein, Anima Anandkumar, and Bryan Catanzaro. Long-short transformer: Efficient transformers for language and vision. Advances in Neural Information Processing Systems, 34:17723-17736, 2021.
[92] Yukun Zhu, Ryan Kiros, Rich Zemel, Ruslan Salakhutdinov, Raquel Urtasun, Antonio Torralba, and Sanja Fidler. Aligning books and movies: Towards story-like visual explanations by watching movies and reading books. In The IEEE International Conference on Computer Vision (ICCV), December 2015.
D.Y.F. Conceptualized the research; coordinated collaborations; developed M2 architectures; led experimental and implementation efforts; assisted in development of theoretical results; coordinated writing.
S.A. Worked on developing M2 architectures; worked on implementing and conducting BERT experiments; worked on implementing and performing CPU experiments; assisted in writing and framing the work.
J.G. Led development of theory and causal algorithms; wrote Appendix D.
I.J. Led development of theory and causal algorithms; wrote Appendix D.
S.E. Worked on writing and framing the work; assisted in development of M2 architectures; assisted in the optimized M2 implementation; conducted mixer benchmarks; assisted with BERT experiments; conducted Swin ImageNet experiments.
A.W.T. Conducted ViT experiments; assisted in writing.
B.S. Assisted in optimized M2 implementation; conducted mixer benchmarks; assisted in writing.
M.P. Assisted in development of M2-GPT architecture.
A.R. Supervised theory development; developed proofs; reviewed manuscript.
C.R. Supervised research; reviewed manuscript.
Simran Arora, Jessica Grogan, Isys Johnson, Sabri Eyuboglu, and Armin Thomas contributed equally to this work.
Appendix A discusses broader impacts of our work. Appendix B presents additional experiments. Appendix C gives details for the experiments, including model architectures and hyperparameters. Appendix D gives missing details and proofs for the theoretical analysis, as well as generalizations to broader results. Appendix E gives a PyTorch code listing of an M2 layer.