JAX PathwaysAll career routesCareer learning guide
Research engineer
Turn a proposed method into an experiment that can support or refute a claim. Your work connects mathematical assumptions, a simple baseline, independent checks and carefully bounded conclusions.
The work
Reproduce a baseline, change one mechanism, and explain whether the result survives a new seed and a fixed evaluation set.
Preparation
Python, array reasoning and calculus help. Complete gradient checking and controlled experiments before comparing architectures or extending derivative rules.
Your portfolio
A reproducibility report with a baseline, ablation, numerical checks and unresolved limitations.
Skills and evidence
JVPs, VJPs, and custom derivatives
Required route practice
Check derivative linearity, finite differences and higher-order behavior.
Explain the independent check, repeat it under a changed condition and retain the observed result.
Practice project: stages 1, 2.
jaxpr and program inspection
Required route practice
Explain supported primitives and reject unsupported programs.
Explain the independent check, repeat it under a changed condition and retain the observed result.
Practice project: stages 3.
Hand-built HMC/VI diagnostics; optional NumPyro extension
Optional project extension; 4 project preparation lessons lie outside the required route
Compare posterior moments and sampler diagnostics against the exact target; NumPyro remains an optional library extension.
Explain the independent check, repeat it under a changed condition and retain the observed result.
Practice project: stages 2, 3.
MLflow experiment provenance
Optional project extension; 3 project preparation lessons lie outside the required route
Freeze the question, retain failed attempts and interpret complete paired results.
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.
Inspect and verify a differentiable method
A custom derivative with an explained jaxpr
Review your readiness
Explain the assumptions of your method and support a conclusion with correctness checks and an appropriate baseline.
- Which independent calculation checks the mechanism you changed?
- What did you hold fixed in the baseline and ablation, and which seeds changed?
- What alternative explanation remains consistent with your results?
- What information besides a Git commit is needed to reproduce an experiment with changed data?
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.