> Markdown version of [/jobs/ext/2298969-senior-staff-machine-learning-engineer-infrastructure](https://www.wearedevelopers.com/jobs/ext/2298969-senior-staff-machine-learning-engineer-infrastructure). Every page supports `.md` or `Accept: text/markdown`. Links point to the HTML versions so they work for humans too. Agent guide: [/agents.md](https://www.wearedevelopers.com/agents.md). --- # Senior/Staff Machine Learning Engineer, Infrastructure - **Company:** Waymo LLC - **Location:** Mountain View, CA, United States - **Experience:** Expert - **Salary:** $251,000.0 - $310,000.0 - **Contract:** Permanent contract - **Skills:** Business Logic, Profiling, Information Engineering, Programming Tools, Distributed Computing Environment, Machine Learning, Tensorflow, Software Engineering, Reinforcement Learning, Pytorch, Large Language Models, Multi-Agent Systems, Machine Learning Operations - **Published:** August 29, 2026 - **Apply:** https://jobs.localjobnetwork.com/apply/add/85683489/1 ## About the Role * 6+ years of professional software engineering experience, with at least 4 years focused on machine learning infrastructure (scaling, training, optimizing, and deploying large-scale ML systems). * Direct ML programming experience on TPU and GPU hardware using frameworks such as JAX, PyTorch, or TensorFlow. * Proven hands-on experience scaling large models using model parallelism, data parallelism, or distributed training techniques. * Strong understanding of state-of-the-art ML models (e.g., autoregressive transformers) and hands-on proficiency with ML accelerator profiling tools to diagnose bottlenecks. * Demonstrated ability to independently lead ambiguous technical initiatives end-to-end and build robust libraries, pipelines, and developer tooling. * Strong verbal and written communication skills to collaborate effectively across distributed, cross-functional teams. Preferred Qualifications * Practical experience in Reinforcement Learning (RL), Sim2Real transfer, or Robotics. * Experience with distributed ML frameworks and accelerators like GPU/TPU. * Domain familiarity with Autonomous Driving systems, multi-agent simulations, or realistic world modeling. ## Description Waymo is an autonomous driving technology company with the mission to be the world's most trusted driver. Since its start as the Google Self-Driving Car Project in 2009, Waymo has focused on building the Waymo Driver-The World's Most Experienced Driver-to improve access to mobility while saving thousands of lives now lost to traffic crashes. The Waymo Driver powers Waymo's fully autonomous ride-hail service and can also be applied to a range of vehicle platforms and product use cases. The Waymo Driver has provided over ten million rider-only trips, enabled by its experience autonomously driving over 100 million miles on public roads and tens of billions in simulation across 15+ U.S. states. Join our ML Infrastructure engineering team advancing state-of-the-art ultra-realistic multi-agent simulations using foundation models. In this role, you will work at the intersection of ML infrastructure, foundation models, and simulation engineering, with a specific focus on writing high-performance business and simulation logic in JAX/TensorFlow running directly on TPUs to power realistic environments for Reinforcement Learning (RL). What You'll Do * Design, build, and optimize realistic simulation environments and business logic running on TPUs using JAX and TensorFlow. Implement and optimize large-scale model and data parallelism strategies for training and running foundation models on TPU hardware. * Collaborate closely with modeling teams to integrate foundation models into simulation pipelines. * Drive technical architectures and system designs from data engineering through simulation execution to meet business and performance objectives. * Profile systems, identify performance bottlenecks across ML accelerators, and optimize end-to-end execution speed. * Translate product and business goals into concrete technical requirements and system deliverables. ## Related Videos - [How Machine Learning is turning the Automotive Industry upside down](https://www.wearedevelopers.com/videos/61-how-machine-learning-is-turning-the-automotive-industry-upside-down) - [Profiling Symfony & PHP apps with Blackfire](https://www.wearedevelopers.com/videos/265-profiling-symfony-php-apps-with-blackfire) - [Machine learning in the browser with TensorFlowjs](https://www.wearedevelopers.com/videos/155-machine-learning-in-the-browser-with-tensorflowjs) - [Photonic Computing: Programming a New Class of AI Accelerators (incl. Live Coding)](https://www.wearedevelopers.com/videos/100196-photonic-computing-programming-a-new-class-of-ai-accelerators-incl-live-coding) - [Getting Started with Machine Learning](https://www.wearedevelopers.com/videos/260-getting-started-with-machine-learning) - [Developing an AI.SDK](https://www.wearedevelopers.com/videos/198-developing-an-ai-sdk) ## Related Articles - [What Are Large Language Models?](https://www.wearedevelopers.com/magazine/304-what-are-large-language-models) - [How to Become an AI Engineer](https://www.wearedevelopers.com/magazine/331-how-to-become-an-ai-engineer) - [MLOps And AI Driven Development](https://www.wearedevelopers.com/magazine/82-mlops-and-ai-driven-development) - [Highest Paying Tech Companies for Developers](https://www.wearedevelopers.com/magazine/220-highest-paying-tech-companies-for-developers) - [Got AI ideas but no money? Here are 10 free ways to level up your AI skills with Google Cloud](https://www.wearedevelopers.com/magazine/600-got-ai-ideas-but-no-money-here-are-10-free-ways-to-level-up-your-ai-skills-with-google-cloud) - [MLops – Deploying, Maintaining And Evolving Machine Learning Models in Production](https://www.wearedevelopers.com/magazine/115-mlops-deploying-maintaining-and-evolving-machine-learning-models-in-production)