JAX PathwaysAll career routesCareer learning guide
Scientific ML engineer
Connect observations to a numerical or physical model. You check derivatives and solver behavior, fit parameters and explain how noise, conditioning and model assumptions affect the answer.
The work
A parameter fit looks accurate but changes sharply when observations move slightly. Separate numerical error from weak identifiability.
Preparation
Bring a physical or numerical question and basic calculus. The shared foundations introduce JAX; the science sequence develops discretization, sensitivity and inverse-problem checks on small CPU examples.
Your portfolio
An inverse-problem report with reference solutions, sensitivity checks and a held-out observation test.
Skills and evidence
scan and batched simulation
Required route practice
Verify batch/time axis contracts and a known trajectory.
Explain the independent check, repeat it under a changed condition and retain the observed result.
Practice project: stages 1.
Solver sensitivities; optional Diffrax extension
Required route practice
Compare the discrete derivative and a refined-grid continuous reference; Diffrax is optional extension work.
Explain the independent check, repeat it under a changed condition and retain the observed result.
Practice project: stages 2.
Numerical accuracy and inverse problems
Required route practice
Separate parameter error, prediction error and solver bias.
Explain the independent check, repeat it under a changed condition and retain the observed result.
Practice project: stages 3.
MLflow experiment provenance
Optional project extension; 3 project preparation lessons lie outside the required route
Attach configuration, source/data identities and a justified comparison to a research notebook.
Explain the independent check, repeat it under a changed condition and retain the observed result.
Practice project: stages 1, 2, 3, 4.
Milestones
Build the shared foundation
A reproducible regression fit with checked gradients and explicit random state.
Recover a hidden physical parameter
A differentiable inverse problem
Review your readiness
Verify parameter recovery, gradient accuracy, and simulation stability on a stated test problem.
- Does the inferred parameter converge when you refine the time grid?
- Do solver sensitivities agree with independent finite differences?
- Can different parameters explain the same observations, and what extra measurement would distinguish them?
- Which model and solver configuration does the recorded experiment actually validate?
This is a learning route for practicing the work, not a claim of job readiness. Use the milestone coverage and a reviewer’s feedback to identify what remains to demonstrate. The capability map below separates required route practice from optional extensions. Public checks and self-reported completion are not independent review or hardware qualification.