Adam O’Brien, PhD

obrienadam89@gmail.com | (603) 266-7012 | Mountain View, CA

GitHub | LinkedIn

Summary

Versatile and impact-driven Software Engineer with an extensive track record in high-performance computing, machine learning infrastructure, data pipelines, and scalable simulation engines. Combines a PhD-level foundation in numerical computing with robust industrial experience building core runtime components, high-throughput data ingestion pipelines, and performance-critical systems natively in C++ and Python.

Work Experience

Google LLC Apr. 2023 – Present

Software Engineer – Machine Learning Infrastructure | Sunnyvale, CA

• Data Infrastructure (PyGrain & JAX Data):
  • Maintain and scale PyGrain and JAX Data, the foundational data pipeline libraries handling ingestion and processing for JAX-based training workloads globally across Google.
  • Own shared platform tooling utilized across deep learning teams to guarantee end-to-end data pipeline reliability, performance optimization, and seamless high-throughput training inputs.
  • Architect and support next-generation serving and inference infrastructure, bridging the gap between massive training-data workflows and low-latency production model delivery.
• ML DNA: Scaling & Infrastructure:
  • Designed and built Google's high-performance JAX CPU embedding library from the ground up, providing specialized infrastructure for accelerated embedding lookups and processing.
  • Engineered JAX Parameter Server Training (PST) and production scaling architectures optimized for massive-scale recommendation and ranking models powering Ads, YouTube, Play, and Geo.
  • Optimized the JAX parameter server infrastructure to deliver up to a 3x throughput improvement in dense worker/server topographies by implementing zero-copy RPC handlers and low-overhead threading patterns.
  • Developed a scalable native C++ optimizer library from scratch, introducing horizontally scalable optimization routines tailored specifically for parameter server topologies.
  • Partnered directly with key Google product areas to guide architectural migration strategies to TPU hardware, maximizing resource utilization for exceptionally large embedding models.
  • Honored with the organization's "Velocity" Award for cross-functional impact in successfully launching JAX PST into production.
• Megascale XLA:
  • Contributed as a core infrastructure engineer on the Megascale XLA team, focusing on the vertical and horizontal scaling challenges of large language and foundation model workloads.
  • Co-authored network topology-aware collective primitives to facilitate efficient distributed ML training clusters operating across cross-metro infrastructure, receiving the internal "Gold Perfy" Award.
  • Designed host-offloaded compute strategies to unblock training efficiency boundaries across multi-slice TPU deployments.

Aurora Innovation Oct. 2021 – Apr. 2023

Senior Software Engineer – Motion Planning Simulation | Mountain View, CA

• Architected core software components for a highly distributed, massive-scale motion planning simulation engine executing millions of autonomous vehicle scenarios daily.
• Developed dynamic agent behavior routing frameworks, introducing adversarial interaction profiles and real-time lane-changing mechanics to evaluate safety-critical edge cases.
• Created automated validation tools for Foreign Object Debris (FOD) simulations, engineering workflows that dynamically combined real-world sensor logs with synthetic environmental augmentations.
• Discovered and executed sweeping runtime optimizations within the mapping and routing engine, generating audited infrastructure savings of ~$130,000 per month in cloud compute expenditures.

Siemens PLM Software Jul. 2019 – Oct. 2021

Advanced Software Engineer – Physics Solvers | Lebanon, NH

• Maintained and optimized parallel fluid dynamics solvers within the flagship industrial simulation suite STAR-CCM+.
• Researched and implemented highly parallelized computational methods to resolve large-scale, complex partial differential equations on multi-node cluster topographies.
• Accelerated core Algebraic Multi-Grid (AMG) linear solvers via high-performance GPU offloading using CUDA.

Education

University of Toronto Sep. 2014 – Jul. 2019

PhD in Computational Science / Applied Mathematics | Toronto, ON, Canada

• Thesis: Development of an immersed boundary method for moving bodies at fluid-fluid-interfaces.
• Authored multiple peer-reviewed publications in elite numerical and computational journals (~60 academic citations).

Technical Skills