ML Infra Engineer (TPU/Jax/Optimization)
Physical IntelligenceSan Francisco, California, United States
Posted January 23, 2026
Full Time
About the Company
Physical Intelligence specializes in integrating general-purpose artificial intelligence with physical systems, creating foundation models and learning algorithms that enhance robotic and actuated device performance. The company’s multidisciplinary team of engineers, scientists, and roboticists develops scalable AI solutions that enable robots to adapt to complex, real-world environments. By focusing on robust, reusable models rather than task-specific code, Physical Intelligence delivers faster deployment and lower maintenance for industrial automation, logistics, and consumer robotics. Their collaborative culture fosters rapid prototyping and continuous learning, positioning them as a leader in the next wave of intelligent physical systems.
Job Description
Help scale and optimize training systems and core model code, building infrastructure for large-scale training and working closely with researchers and model engineers.
Requirements
- Strong software engineering fundamentals
- Experience building ML training infrastructure or internal platforms
- Hands-on large-scale training experience in JAX (preferred), PyTorch
- Familiarity with distributed training, multi-host setups, data loaders, and evaluation pipelines
- Experience managing training workloads on cloud platforms (e.g., SLURM, Kubernetes, GCP TPU/GKE, AWS)
- Ability to debug and optimize performance bottlenecks across the training stack
- Strong cross-functional communication and ownership mindset
