This job post has expired on August 15, 2026. It is likely that the position has already been filled.
JAX Internals Expert at AfterQuery
posted 3 months agoJAX Internals Expert | $75/hr | Worldwide Remote
AFTERQUERY is seeking highly skilled ML, systems, and distributed systems engineers with deep, internals-level JAX expertise. This is a contractor engagement ideal for professionals with hands-on experience in JAX's lower-level mechanics, including compiler work, custom transformations, and distributed execution.
What You'll Be Working On
You'll be contributing at the internals level of JAX-based systems, working across areas such as custom transformations, XLA/MLIR compiler integration, distributed sharding strategies, and custom kernel development.
Required Qualifications
- Full-time professional or research experience working with JAX
- Demonstrated internals-level expertise in at least one of the following:
- Custom transformations: grad, vmap, pmap, custom_vjp, custom_jvp
- Custom primitives and lowering rules
- XLA/MLIR compiler-level development
- Distributed sharding: pjit, jax.Array, GSPMD
- Custom kernel authoring via Pallas
- Access to suitable hardware — GPU/TPU-enabled machine or cloud instance required
Compensation
$75/hr | Department: Software Engineering
Surface-level JAX experience will not be considered — deep, internals-level expertise is required to qualify.
