JAX PathwaysML & training engineer
Build models that train reliably.
- Build and evaluate neural networks
- Design data pipelines and resumable training
HomeCareer paths
Explore JAX careers.
Explore the problems each role solves, the evidence you can practice producing and the lessons available now.
Each route separates available lessons from planned work.
Research engineer
Turn mathematical ideas into experiments.
- Prototype differentiable algorithms
- Investigate uncertainty and inference
Scientific ML engineer
Simulate systems. Learn their hidden parameters.
- Build and batch numerical simulations
- Differentiate through dynamics and solvers
RL engineer
Build agents that learn by interacting.
- Build reproducible environments and rollouts
- Implement and diagnose policy updates
Accelerator & performance engineer
Find bottlenecks. Scale the right computation.
- Diagnose compute, input, and communication bottlenecks
- Partition workloads across devices
Inference & deployment engineer
Make a trained model usable.
- Adapt and export trained computations
- Measure serving latency and throughput
JAX / accelerator operations engineer
Keep real workloads observable and recoverable.
- Manage accelerator jobs and runtime lifecycle
- Monitor throughput, stalls, and checkpoint health