Ajay Jain

Doctoral student (Ph.D.)

Berkeley Artificial Intelligence Research (BAIR)
Department of Electrical Engineering and Computer Science (EECS)
UC Berkeley

Email: ajayj at berkeley dot edu


I'm a PhD student at UC Berkeley, affiliated with Berkeley Artificial Intelligence Research. I work on machine learning and combinatorial optimization, with a special interest in systems applications. I graduated from MIT with an S.B. in Computer Science in 2019, where I was advised by Saman Amarasinghe and affiliated with the MIT COMMIT lab.

If you go to MIT, you've likely received an email from me at some point about a machine learning reading group or talk -- from 2017-2018, I was the co-president of MIT MIC, and volunteered my time by ordering a tremendous amount of pizza to entice students to read papers together. Now, I sit on the board of the non-profit Machine Intelligence Community. In the past, I've organized HackMIT and Blueprint.


New! June 2019: Presented a 10 minute talk at ISCA 2019 Workshop on Machine Learning for Systems
New! June 2019: Poster at ICML 2019 Workshop on Phenomena of Deep Learning
New! June 2019: Started my PhD at UC Berkeley
Dec 2018: Revec accepted to Compiler Construction 2019
June 2018: Headed to Toronto for a research internship at Uber Advanced Technologies Group
June 2017: Joining Facebook's Applied Machine Learning team for the summer
August 2016: Started my undergrad at MIT
June 2016: Interning at Juniper Networks this summer

Publications * indicates equal contribution

Revec: Program Rejuvenation through Revectorization

Charith Mendis*, Ajay Jain*, Paras Jain, Saman Amarasinghe
Proceedings of the 28th International Conference on Compiler Construction (CC), Feb 2019

Autonomy for Surface Ship Interception

C. Mirabito, D.N. Subramani, T. Lolla, J. P.J. Haley, A. Jain, P.F.J. Lermusiaux, C. Li, D. Yue, Y. Liu, F. Hover, N. Pulsone, J. Edwards, K. Railey, and G. Shaw
Proceedings of the 60th OCEANS Conference, Oceans MTS/IEEE Aberdeen, June 2017

Workshop papers

Learning Automatic Schedulers with Projective Reparameterization

Ajay Jain, Saman Amarasinghe
The 46th International Symposium on Computer Architecture (ISCA)
Workshop on Machine Learning for Systems, June 2019
10 minute talk

Using effective dimension to analyze feature transformations in deep neural networks

Kavya Ravichandran, Ajay Jain, Alexsander Rakhlin
The 36th International Conference on Machine Learning (ICML)
Workshop on Identifying and Understanding Deep Learning Phenomena, June 2019

Dynamic Space-Time scheduling for GPU inference

Paras Jain, Xiangxi Mo, Ajay Jain, Harikaran Subbaraj, Rehan Sohail Durrani, Alexey Tumanov, Joseph Gonzales, and Ion Stoica
The 32nd Annual Conference on Neural Information Processing Systems (NeurIPS)
Workshop on Systems for Machine Learning, Dec 2018


The OoO VLIW JIT Compiler for GPU Inference

Paras Jain, Xiangxi Mo, Ajay Jain, Alexey Tumanov, Joseph E Gonzalez, Ion Stoica
arXiv:1901.10008 preprint, Jan 2019