AWS Neuron is the complete software stack for the Inferentia and Trainium cloud-scale machine learning accelerators. As a part of the Neuron Frameworks team you'll deliver on PyTorch and JAX for AWS Neuron, working with the open source ecosystem.
You will develop and improve ML frameworks to deliver an outstanding user experience for PyTorch and JAX ML model developers. You will work closely with teams across AWS Neuron including compiler, training and inference optimization to optimize frameworks for AWS's accelerator architectures, and engage closely with the PyTorch and JAX and other ML Framework communities to take advantage of their latest capabilities and improve performance and usability.
A successful candidate will have experience developing Machine Learning infrastructure and/or ML Frameworks on XLA devices and corresponding framework components such as Torch-XLA, Torch Dynamo, Inductor, Open-XLA project integrations using PJRT or StableHLO, alongwith familarity of OpenXLA compilers. The ideal candidate should have a demonstrated ability to work with open source communities to influence future community direction, a solid software development experience and a motivation to achieve results. Knowledge of optimization of transformers is highly valued.
Key job responsibilities
As you develop and enhance features for ML Frameworks you’ll build high-impact solutions with and for users at Amazon and across the ML industry. You will and work with the JAX and PyTorch open source communities to adopt and develop new technologies, and contribute innovations.
You'll collaborate with other engineers within Neuron and cross-functionally to help drive projects to successful business outcomes using your technical skills.
You will work in a fast-paced development environment, where you’re always working on the most important challenges and customer needs.