Anjiang Wei

I'm Anjiang Wei (魏安江 in Chinese), a research scientist at Google DeepMind. I obtained my PhD in Computer Science from Stanford University, where I worked under the guidance of Alex Aiken. I obtained my B.S. in Computer Science from Turing Class, Peking University.

My research interest is LLMs for challenges in software systems, especially performance and correctness.

Email  /  Scholar  /  Twitter  /  Github

Research

• AccelOpt: A Self-Improving LLM Agentic System for AI Accelerator Kernel Optimization
  Genghan Zhang, Shaowei Zhu, Anjiang Wei, Zhenyu Song, Allen Nie, Zhen Jia, Nandita Vijaykumar, Yida Wang, Kunle Olukotun
  Preprint
  

• Equivalence Checking of ML GPU Kernels
  Kshitij Dubey, Benjamin Driscoll, Anjiang Wei, Neeraj Kayal, Rahul Sharma, Alex Aiken
  Preprint
  

• Improving Parallel Program Performance with LLM Optimizers via Agent-System Interfaces
  Anjiang Wei*, Allen Nie*, Thiago S. F. X. Teixeira, Rohan Yadav, Wonchan Lee, Ke Wang, Alex Aiken
  International Conference on Machine Learning
  Oral Presentation at PRAL Workshop
  ICML 2025, Vancouver, Canada, July 2025
  slides / poster / video

• CodeARC: Benchmarking Reasoning Capabilities of LLM Agents for Inductive Program Synthesis
  Anjiang Wei, Tarun Suresh, Jiannan Cao, Naveen Kannan, Yuheng Wu, Kai Yan, Thiago S. F. X. Teixeira, Ke Wang, Alex Aiken
  Conference on Language Modeling
  COLM 2025, Montreal, Canada, October 2025
  code / dataset / website

• SATBench: Benchmarking LLMs' Logical Reasoning via Automated Puzzle Generation from SAT Formulas
  Anjiang Wei, Yuheng Wu, Yingjia Wan, Tarun Suresh, Huanmi Tan, Zhanke Zhou, Sanmi Koyejo, Ke Wang, Alex Aiken
  Conference on Empirical Methods in Natural Language Processing
  EMNLP 2025, Suzhou, China, November 2025
  code / dataset / website

• EquiBench: Benchmarking Large Language Models' Reasoning about Program Semantics via Equivalence Checking
  Anjiang Wei, Jiannan Cao, Ran Li, Hongyu Chen, Yuhui Zhang, Ziheng Wang, Yuan Liu, Thiago S. F. X. Teixeira, Diyi Yang, Ke Wang, Alex Aiken
  Conference on Empirical Methods in Natural Language Processing
  EMNLP 2025, Suzhou, China, November 2025
  code / dataset / website

• SuperCoder: Assembly Program Superoptimization with Large Language Models
  Anjiang Wei, Tarun Suresh, Huanmi Tan, Yinglun Xu, Gagandeep Singh, Ke Wang, Alex Aiken
  The 4th Deep Learning for Code Workshop
  DL4C workshop at NeurIPS 2025, San Diego, California, December 2025
  dataset

• Astra: A Multi-Agent System for GPU Kernel Performance Optimization
  Anjiang Wei, Tianran Sun, Yogesh Seenichamy, Hang Song, Anne Ouyang, Azalia Mirhoseini, Ke Wang, Alex Aiken
  The 4th Deep Learning for Code Workshop
  DL4C workshop at NeurIPS 2025, San Diego, California, December 2025
  code

• Quokka: Accelerating Program Verification with LLMs via Invariant Synthesis
  Anjiang Wei, Tarun Suresh, Tianran Sun, Haoze Wu, Ke Wang, Alex Aiken
  Preprint
  dataset

• VeriCoder: Enhancing LLM-Based RTL Code Generation through Functional Correctness Validation
  Anjiang Wei, Huanmi Tan, Tarun Suresh, Daniel Mendoza, Thiago S. F. X. Teixeira, Ke Wang, Caroline Trippel, Alex Aiken
  The 4th Deep Learning for Code Workshop
  DL4C workshop at NeurIPS 2025, San Diego, California, December 2025
  code / dataset / model / website

• Mapple: A Domain-Specific Language for Mapping Distributed Programs
  Anjiang Wei, Rohan Yadav, Hang Song, Wonchan Lee, Ke Wang, Alex Aiken
  Preprint
  code

• Task-Based Programming for Adaptive Mesh Refinement in Compressible Flow Simulations
  Anjiang Wei, Hang Song, Mert Hidayetoglu, Elliott Slaughter, Sanjiva K. Lele, Alex Aiken
  Preprint
  code

• Automated Generation of Challenging Multiple-Choice Questions for Vision Language Model Evaluation
  Yuhui Zhang, Yuchang Su, Yiming Liu, Xiaohan Wang, James Burgess, Elaine Sui, Chenyu Wang, Josiah Aklilu, Alejandro Lozano, Anjiang Wei, Ludwig Schmidt, Serena Yeung-Levy
  IEEE / CVF Computer Vision and Pattern Recognition Conference
  CVPR 2025, Nashville, Tennessee, June 2025
  website / code

• PBA: Percentile-Based Level Allocation for Multiple-Bits-Per-Cell RRAM
  Anjiang Wei, Akash Levy, Pu (Luke) Yi, Robert M. Radway, Priyanka Raina, Subhasish Mitra, Sara Achour
  International Conference on Computer Aided Design
  ICCAD 2023, San Francisco, CA, October 2023
  slides

• Fuzzing Deep-Learning Libraries via Automated Relational API Inference
  Yinlin Deng*, Chenyuan Yang*, Anjiang Wei, Lingming Zhang
  ACM Joint European Software Engineering Conference and Symposium on the Foundations of Software Engineering
  ESEC/FSE 2022, Singapore, November 2022
  code

• AMOS: Enabling Automatic Mapping for Tensor Computations On Spatial Accelerators with Hardware Abstraction
  Size Zheng, Renze Chen, Anjiang Wei, Yicheng Jin, Qin Han, Liqiang Lu, Bingyang Wu, Xiuhong Li, Shengen Yan, Yun Liang
  International Symposium on Computer Architecture
  ISCA 2022, New York City, NY, June 2022
  code / slides

• Free Lunch for Testing: Fuzzing Deep-Learning Libraries from Open Source
  Anjiang Wei, Yinlin Deng, Chenyuan Yang, Lingming Zhang
  International Conference on Software Engineering
  ICSE 2022, Pittsburgh, PA, May 2022
  code / slides / video

• Preempting Flaky Tests via Non-Idempotent-Outcome Tests
  Anjiang Wei, Pu Yi, Zhengxi Li, Tao Xie, Darko Marinov, Wing Lam
  International Conference on Software Engineering
  ICSE 2022, Pittsburgh, PA, May 2022
  dataset / slides / video

• NeoFlow: A Flexible Framework for Enabling Efficient Compilation for High Performance DNN Training
  Size Zheng, Renze Chen, Yicheng Jin, Anjiang Wei, Bingyang Wu, Xiuhong Li, Shengen Yan, Yun Liang
  IEEE Transactions on Parallel and Distributed Systems
  TPDS 2022

• Finding Polluter Tests Using Java PathFinder
  Pu Yi, Anjiang Wei, Wing Lam, Tao Xie, Darko Marinov
  ACM SIGSOFT Software Engineering Notes
  SEN 2021
  slides

• Domain-Specific Fixes for Flaky Tests with Wrong Assumptions on Underdetermined Specifications
  Peilun Zhang, Yanjie Jiang, Anjiang Wei, Victoria Stodden, Darko Marinov, August Shi
  International Conference on Software Engineering
  ICSE 2021, Virtual Conference, May 2021
  slides / video

• Probabilistic and Systematic Coverage of Consecutive Test-Method Pairs for Detecting Order-Dependent Flaky Tests
  Anjiang Wei, Pu Yi, Tao Xie, Darko Marinov, Wing Lam
  International Conference on Tools and Algorithms for the Construction and Analysis of Systems
  TACAS 2021, Virtual Conference, March 2021
  slides / video

• A Large-Scale Longitudinal Study of Flaky Tests
  Wing Lam, Stefan Winter, Anjiang Wei, Tao Xie, Darko Marinov, Jonathan Bell
  ACM SIGPLAN International Conference on Object-Oriented Programming, Systems, Languages, and Applications
  OOPSLA 2020, Virtual Conference, November 2020
  slides / video

Teaching

CS161 - Design and Analysis of Algorithms (Summer 2024)
CS224N - Natural Language Processing with Deep Learning (Winter 2025)

Miscellanea

I previously organized the Software Research Lunch (subscribe) and the Stanford Systems Seminar (subscribe).

I enjoy playing piano (performed at the Shanghai Concert Hall), chess, learning physics (Top 10 in Shanghai in the National Physics Olympiad), and ball games (ping pong, badminton, foosball).

Template taken from Jon Barron.