Startup seeking a senior freelance JAX/Scientific Computing Engineer to implement a locked mathematical specification in production JAX. Math already designed by an applied mathematician. Remote, 30–35 hrs/week, long‑term contract. NDA required before detailed spec access.
Technical surface (you don’t need all of it on day one, but you’ll touch most of it):
- High‑dimensional square‑root unscented Kalman filtering (Joseph‑form updates, adaptive noise)
- SDE integration (Euler–Maruyama, Milstein, exact‑transition schemes)
- Hierarchical Bayesian inference via NumPyro NUTS (running pre‑specified models, not designing them)
- Amortized neural posterior estimation via conditional normalizing flows (BayesFlow)
- Multi‑state hazard modelling, doubly‑robust estimators, constrained contextual bandits
- JAX → ONNX export pipeline (jax2tf → tf2onnx, golden‑set parity testing)
Stack:JAX, Equinox, Diffrax, NumPyro, BayesFlow, Triton Inference Server, vLLM (co‑tenant on GPU node), Dagster, FastAPI, ONNX Runtime (on edge devices). Single H100‑class GPU, scaling planned.
What we need from you:
- Production JAX experience (`vmap`, `lax.scan`, JIT compilation, FP8/BF16 memory budgeting)
- Familiarity with Equinox or Flax, Diffrax (or another ODE/SDE library), NumPyro
- At least one model‑export pipeline (jax2tf + tf2onnx, or equivalent)
- Ability to read a math spec and turn it into efficient, tested JAX code without debating the upstream design choices
What this is not:
- Not a math‑design role (priors, estimators, and model structure are locked)
- Not a PyTorch job repackaged — we need genuine JAX‑internals fluency
Engagement:Hourly contract, 20-30 hrs/week, 4‑month initial term with strong intent to continue. Rate $80–100/hr depending on background. Fully remote, async‑first. Weekly 30‑min architecture sync.
Next step:If interested, please reply here or email me at alicharanek@gmail.com with a few lines about your JAX experience. We’ll schedule a 30‑min call, and if it’s a mutual fit, we’ll share the full spec under NDA.
No recruiters, please.