JAX Internals Expert at AfterQuery

JAX 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.