JAX/ONNX/Scientific Computing Engineer Job /Free Lance (contract,remote)

[alicharanek]

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, 20-30 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.