Generalizing Shape Analysis with Gradual Types

Authors Zeina Migeed , James Reed, Jason Ansel , Jens Palsberg



PDF
Thumbnail PDF

File

LIPIcs.ECOOP.2024.29.pdf
  • Filesize: 0.99 MB
  • 28 pages

Document Identifiers

Author Details

Zeina Migeed
  • University of California, Los Angeles (UCLA), CA, USA
James Reed
  • Fireworks AI, Redwood City, CA, USA
Jason Ansel
  • Meta, Menlo Park, CA, USA
Jens Palsberg
  • University of California, Los Angeles (UCLA), CA, USA

Acknowledgements

We thank Akshay Utture, Micky Abir and Shuyang Liu for helpful comments.

Cite AsGet BibTex

Zeina Migeed, James Reed, Jason Ansel, and Jens Palsberg. Generalizing Shape Analysis with Gradual Types. In 38th European Conference on Object-Oriented Programming (ECOOP 2024). Leibniz International Proceedings in Informatics (LIPIcs), Volume 313, pp. 29:1-29:28, Schloss Dagstuhl – Leibniz-Zentrum für Informatik (2024)
https://doi.org/10.4230/LIPIcs.ECOOP.2024.29

Abstract

Tensors are multi-dimensional data structures that can represent the data processed by machine learning tasks. Tensor programs tend to be short and readable, and they can leverage libraries and frameworks such as TensorFlow and PyTorch, as well as modern hardware such as GPUs and TPUs. However, tensor programs also tend to obscure shape information, which can cause shape errors that are difficult to find. Such shape errors can be avoided by a combination of shape annotations and shape analysis, but such annotations are burdensome to come up with manually. In this paper, we use gradual typing to reduce the barrier of entry. Gradual typing offers a way to incrementally introduce type annotations into programs. From there, we focus on tool support for type migration, which is a concept that closely models code-annotation tasks and allows us to do shape reasoning and utilize it for different purposes. We develop a comprehensive gradual typing theory to reason about tensor shapes. We then ask three fundamental questions about a gradually typed tensor program. (1) Does the program have a static migration? (2) Given a program and some arithmetic constraints on shapes, can we migrate the program according to the constraints? (3) Can we eliminate branches that depend on shapes? We develop novel tools to address the three problems. For the third problem, there are currently two PyTorch tools that aim to eliminate branches. They do so by eliminating them for just a single input. Our tool is the first to eliminate branches for an infinite class of inputs, using static shape information. Our tools help prevent bugs, alleviate the burden on the programmer of annotating the program, and improves the process of program transformation.

Subject Classification

ACM Subject Classification
  • Software and its engineering → Software notations and tools
Keywords
  • Tensor Shapes
  • Gradual Types
  • Migration

Metrics

  • Access Statistics
  • Total Accesses (updated on a weekly basis)
    0
    PDF Downloads

References

  1. Martín Abadi, Ashish Agarwal, Paul Barham, Eugene Brevdo, Zhifeng Chen, Craig Citro, Greg S. Corrado, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, Ian Goodfellow, Andrew Harp, Geoffrey Irving, Michael Isard, Yangqing Jia, Rafal Jozefowicz, Lukasz Kaiser, Manjunath Kudlur, Josh Levenberg, Dan Mane, Rajat Monga, Sherry Moore, Derek Murray, Chris Olah, Mike Schuster, Jonathon Shlens, Benoit Steiner, Ilya Sutskever, Kunal Talwar, Paul Tucker, Vincent Vanhoucke, Vijay Vasudevan, Fernanda Viegas, Oriol Vinyals, Pete Warden, Martin Wattenberg, Martin Wicke, Yuan Yu, and Xiaoqiang Zheng. Tensorflow: Large-scale machine learning on heterogeneous distributed systems, 2016. URL: https://doi.org/10.48550/arXiv.1603.04467.
  2. Jason Ansel. TorchDynamo. Software release, January 2022. URL: https://github.com/pytorch/torchdynamo.
  3. Jason Ansel, Animesh Jain, David Berard, Will Constable, Will Feng, Sherlock Huang, Mario Lezcano, CK Luk, Matthias Reso, Michael Suo, William Wen, Richard Zou, Edward Yang, Michael Voznesensky, Evgeni Burovski, Alban Desmaison, Jiong Gong, Kshiteej Kalambarkar, Yanbo Liang, Bert Maher, Mark Saroufim, Phil Tillet, Shunting Zhang, Ajit Mathews, Horace He, Bin Bao, Geeta Chauhan, Zachary DeVito, Michael Gschwind, Laurent Kirsch, Jason Liang, Yunjie Pan, Marcos Yukio Siraichi, Eikan Wang, Xu Zhao, Gregory Chanan, Natalia Gimelshein, Peter Bell, Anjali Chourdia, Elias Ellison, Brian Hirsh, Michael Lazos, Yinghai Lu, Christian Puhrsch, Helen Suk, Xiaodong Wang, Keren Zhou, Peng Wu, and Soumith Chintala. Pytorch 2: Faster machine learning through dynamic Python bytecode transformation and graph compilation. In Proceedings of ASPLOS'24, International Conference on Architectural Support for Programming Languages and Operating Systems, 2024. Google Scholar
  4. Gavin Bierman, Martín Abadi, and Mads Torgersen. Understanding typescript. In Richard Jones, editor, ECOOP 2014 - Object-Oriented Programming, pages 257-281, Berlin, Heidelberg, 2014. Springer Berlin Heidelberg. Google Scholar
  5. James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milne, and Qiao Zhang. JAX: composable transformations of Python+NumPy programs. Software release, 2018. URL: http://github.com/google/jax.
  6. Matteo Cimini and Jeremy Siek. The gradualizer: A methodology and algorithm for generating gradual type systems. In Proceedings of POPL'16, ACM Symposium on Principles of Programming Languages, New York, 2016. ACM. Google Scholar
  7. Ben Greenman and Matthias Felleisen. A spectrum of type soundness and performance. Proc. ACM Program. Lang., 2(ICFP), July 2018. URL: https://doi.org/10.1145/3236766.
  8. Ben Greenman and Zeina Migeed. On the cost of type-tag soundness. In PEPM, 2018. Google Scholar
  9. Momoko Hattori, Naoki Kobayashi, and Ryosuke Sato. Gradual tensor shape checking, 2022. URL: https://doi.org/10.48550/arXiv.2203.08402.
  10. Ho Young Jhoo, Sehoon Kim, Woosung Song, Kyuyeon Park, DongKwon Lee, and Kwangkeun Yi. A static analyzer for detecting tensor shape errors in deep neural network training code. In Proceedings of the ACM/IEEE 44th International Conference on Software Engineering (ICSE), 2022. Google Scholar
  11. Sifis Lagouvardos, Julian T Dolby, Neville Grech, Anastasios Antoniadis, and Yannis Smaragdakis. Static analysis of shape in tensorflow programs. In ECOOP, Germany, 2020. LIPICS. Google Scholar
  12. Zeina Migeed and Jens Palsberg. What is decidable about gradual types? Proc. ACM Program. Lang., 4(POPL), December 2019. URL: https://doi.org/10.1145/3371097.
  13. Adam Paszke, Sam Gross, Soumith Chintala, Gregory Chanan, Edward Yang, Zachary DeVito, Zeming Lin, Alban Desmaison, Luca Antiga, and Adam Lerer. Automatic differentiation in pytorch. In NIPS 2017 Workshop on Autodiff, 2017. URL: https://openreview.net/forum?id=BJJsrmfCZ.
  14. Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang, Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. PyTorch: An Imperative Style, High-Performance Deep Learning Library. In Advances in Neural Information Processing Systems 32, pages 8024-8035. Curran Associates, Inc., 2019. URL: http://papers.neurips.cc/paper/9015-pytorch-an-imperative-style-high-performance-deep-learning-library.pdf.
  15. Adam Paszke and Brennan Saeta. Tensors fitting perfectly, 2021. URL: https://doi.org/10.48550/arXiv.2102.13254.
  16. Luna Phipps-Costin, Carolyn Jane Anderson, Michael Greenberg, and Arjun Guha. Solver-based gradual type migration. In ACM SIGPLAN Conference on Object Oriented Programming, Systems, Languages and Applications (OOPSLA), New York, 2021. ACM. Google Scholar
  17. James K. Reed, Zachary DeVito, Horace He, Ansley Ussery, and Jason Ansel. Torch.fx: Practical program capture and transformation for deep learning in python. Accessed Jul 12, 2024, 2021. URL: https://doi.org/10.48550/arXiv.2112.08429.
  18. Norman A. Rink. Modeling of languages for tensor manipulation. ArXiv, abs/1801.08771, 2018. Google Scholar
  19. Jared Roesch, Steven Lyubomirsky, Marisa Kirisame, Josh Pollock, Logan Weber, Ziheng Jiang, Tianqi Chen, Thierry Moreau, and Zachary Tatlock. Relay: A high-level IR for deep learning. CoRR, abs/1904.08368, 2019. URL: https://arxiv.org/abs/1904.08368.
  20. Jared Roesch, Steven Lyubomirsky, Logan Weber, Josh Pollock, Marisa Kirisame, Tianqi Chen, and Zachary Tatlock. Relay: a new IR for machine learning frameworks. In Proceedings of the 2nd ACM SIGPLAN International Workshop on Machine Learning and Programming Languages. ACM, June 2018. URL: https://doi.org/10.1145/3211346.3211348.
  21. Sanjit A. Seshia and Randal E. Bryant. Deciding quantifier-free presburger formulas using parameterized solution bounds. Logical Methods in Computer Science, 1:1-26, 2005. Google Scholar
  22. Jeremy G. Siek and Walid Taha. Gradual typing for functional languages. In IN SCHEME AND FUNCTIONAL PROGRAMMING WORKSHOP, pages 81-92, 2006. Google Scholar
  23. Jeremy G. Siek, Michael M. Vitousek, Matteo Cimini, and John Tang Boyland. Refined criteria for gradual typing. In SNAPL, pages 274-293, Germany, 2015. LIPICS. Google Scholar
  24. Chengnian Sun, Yuanbo Li, Qirun Zhang, Tianxiao Gu, and Zhendong Su. Perses: Syntax-guided program reduction. In ICSE'18, International Conference on Software Engineering, 2018. Google Scholar
  25. Asumu Takikawa, Daniel Feltey, Ben Greenman, Max S. New, Jan Vitek, and Matthias Felleisen. Is sound gradual typing dead? In Proceedings of the 43rd Annual ACM SIGPLAN-SIGACT Symposium on Principles of Programming Languages, pages 456-468, 2016. Google Scholar
  26. Sam Tobin-Hochstadt and Matthias Felleisen. The design and implementation of typed scheme. In Proceedings of the 35th Annual ACM SIGPLAN-SIGACT Symposium on Principles of Programming Languages, POPL '08, pages 395-406, New York, NY, USA, 2008. ACM. URL: https://doi.org/10.1145/1328438.1328486.
  27. Sahil Verma and Zhendong Su. Shapeflow: Dynamic shape interpreter for tensorflow, 2020. URL: https://doi.org/10.48550/arXiv.2011.13452.
  28. Michael M. Vitousek, Cameron Swords, and Jeremy G. Siek. Big types in little runtime: Open-world soundness and collaborative blame for gradual type systems. SIGPLAN Not., 52(1):762-774, 2017. Google Scholar
  29. Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Perric Cistac, Clara Ma, Yacine Jernite, Julien Plu, Canwen Xu, Teven Le Scao, Sylvain Gugger, Mariama Drame, Quentin Lhoest, and Alexander M. Rush. Transformers: State-of-the-Art Natural Language Processing. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing: System Demonstrations, pages 38-45. Association for Computational Linguistics, October 2020. URL: https://www.aclweb.org/anthology/2020.emnlp-demos.6.
  30. Wayne Wolf. Computers as Components, Principles of Embedded Computing System Design. Morgan Kaufman Publishers, 2000. Google Scholar
  31. Yuhao Zhang, Yifan Chen, Shing-Chi Cheung, Yingfei Xiong, and Lu Zhang. An empirical study on tensorflow program bugs. In Proceedings of the 27th ACM SIGSOFT International Symposium on Software Testing and Analysis, ISSTA 2018, pages 129-140, New York, NY, USA, 2018. Association for Computing Machinery. URL: https://doi.org/10.1145/3213846.3213866.
Questions / Remarks / Feedback
X

Feedback for Dagstuhl Publishing


Thanks for your feedback!

Feedback submitted

Could not send message

Please try again later or send an E-mail