Applied Scientist, Neuron ARG, Annapurna ML

About the role

The Automated Reasoning Group in the Amazon Neuron team is seeking an Applied Scientist to work at the intersection of Artificial Intelligence and program analysis. This role focuses on raising the code quality bar in our state-of-the-art deep learning compiler stack, which optimizes application models across diverse domains, including Large Language and Vision models from frameworks like PyTorch and JAX.

You will work closely with our custom-built Machine Learning accelerator, Trainium, a cutting-edge innovation for advanced ML capabilities and the foundation of Generative AI. Your contributions will include designing, developing, and deploying analyzers for ML compiler stages and compiler IRs, architecting business-critical tooling, publishing research, and mentoring a team of experienced scientists and engineers.

Responsibilities

• Design, develop, and deploy analyzers for ML compiler stages and compiler IRs.
• Architect and implement business-critical tooling for the deep learning compiler stack.
• Publish research and mentor a team of scientists and engineers.
• Tackle critical challenges alongside a talented engineering team to advance compiler technology and deep-learning systems software.
• Innovate as a trusted AWS Neuron engineer on behalf of customers.

Requirements

• PhD, or Master’s degree and 4+ years of experience in Computer Science, Computer Engineering, Machine Learning, or a related field.
• Strong experience in programming languages, compilers, program analyzers, theorem provers, or program synthesis engines.
• Technical credibility and curiosity as an individual contributor.

Preferred Qualifications

• PhD or equivalent research experience.
• Background in machine learning and AI accelerators.

Benefits

• Comprehensive health insurance (medical, dental, vision, prescription, Basic Life & AD&D, with options for supplemental life plans).
• Employee Assistance Program (EAP) and Mental Health Support.
• Medical Advice Line and Flexible Spending Accounts.
• Adoption and Surrogacy Reimbursement coverage.
• 401(k) matching.
• Paid time off and parental leave.
• Sign-on payments and restricted stock units (RSUs).

Learn more about our benefits at amazon.jobs/en/benefits.

Pay

• USA, CA, Cupertino: $171,600.00 - $222,200.00 USD annually.
• USA, WA, Seattle: $142,800.00 - $193,200.00 USD annually.

Final compensation will be determined based on experience, qualifications, and location.