Skip to main content
CareerApp

Skill

JAX

Data Science, Analytics and AI/ML

JAX is a high-performance numerical computing library developed by Google that combines NumPy-like syntax with automatic differentiation and just-in-time compilation via XLA, enabling fast execution on CPUs, GPUs, and TPUs. It is heavily used by machine learning researchers and engineers for building and training neural networks, especially in research settings requiring custom gradients and transformations like vmap and pmap. It underlies libraries such as Flax and Haiku and is popular for large-scale deep learning research.

JAX in the job market

Last checked October 28, 2026

Open roles
42

on Career App

Employers
16

hiring for it

Median pay
$283,000

from 26 roles

Disclose pay
62%

of these roles

Open roles requiring JAX (42)

Research Engineer, Safety & Alignment

Cognition AI

San Francisco, CA

This role leads safety and alignment work for Devin, an autonomous AI agent that writes code and takes actions in production systems. You'll design evaluations, red-team the agent, develop alignment methods, and partner with the agent team to ensure safe behavior at scale—a hands-on research engineering position where safety work directly influences a deployed product.

Listed on Cognition AI’s careers site · Apply there ↗

Staff+ Software Engineer, ML Inference Path

Menlo Ventures Portfolio

Full-time · San Francisco, CA · $320,000 – $485,000

This role builds the production infrastructure that deploys machine learning safety systems at scale across Claude's platforms, working at the intersection of ML engineering, distributed systems, and AI safety to operationalize research into production classifiers and defenses.

Listed on Menlo Ventures Portfolio’s careers site · Apply there ↗

Senior Software Engineer, GPU Infrastructure (HPC)

Cohere

Full-time · Canada, United States · CA$285,000 – CA$340,000 · Remote

This role builds and operates GPU/TPU infrastructure at scale for training frontier AI models, requiring deep expertise in Kubernetes, distributed systems, and close collaboration with AI researchers. It suits engineers with strong systems knowledge who thrive on solving complex infrastructure challenges in fast-paced environments.

Listed on Cohere’s careers site · Apply there ↗

KLA-Tencor

Algorithm Engineer - Deep Learning

KLA-Tencor

Full-time · Milpitas, CA · $136,300 – $231,700

KLA seeks a deep learning engineer to advance computer vision and image processing for semiconductor manufacturing inspection systems. This role involves developing and optimizing state-of-the-art models across the full lifecycle, from conception through deployment, for a leader in process control technology.

Listed on KLA-Tencor’s careers site · Apply there ↗

Target

Sr Applied Data Scientist - Personalization (applied ML, PyTorch, RecSys)

Target

7000 Target Pkwy N,NCD-0375 Brooklyn Park,MN 55445 · $98,000 – $176,000

Target is seeking a senior data scientist to build and optimize recommendation and personalization systems using deep learning and machine learning at scale. This role suits experienced practitioners who want to apply their expertise in PyTorch, ranking algorithms, and large-scale data systems to solve business problems in a collaborative environment.

Listed on Target’s careers site · Apply there ↗

Principal AI Researcher

Workday

Full-time · $228,000 – $342,000

This role leads Workday's newly formed AI Research team in advancing large language models and autonomous agents for enterprise applications. It suits seasoned AI researchers with publications or shipped products who want to define cutting-edge research directions while bridging theory and production-scale impact.

Listed on Workday’s careers site · Apply there ↗

NVIDIA

Engineering Manager, AI Compiler Analysis

NVIDIA

Full-time · $168,000 – $322,000

Lead a technical team building verification systems for NVIDIA's AI compilers, ensuring deep learning models compile correctly to GPU code. This role suits experienced engineering managers who understand compiler optimization and want to drive quality at the intersection of AI frameworks and systems engineering.

Listed on NVIDIA’s careers site · Apply there ↗

AI Robotics Researcher Intern (Dexterous Manipulation)

NIO

Internship · San Jose, CA · $38 – $46 / hour

NIO seeks an AI robotics research intern to develop frameworks for teaching robots dexterous manipulation by learning from human demonstrations and large-scale data, with a focus on translating human skills into real-world robotic behaviors through simulation and hardware deployment.

Listed on NIO’s careers site · Apply there ↗

Expedia Group

Senior Machine Learning Scientist - Agentic Experience

Expedia Group

$173,000 – $242,500

This role leads the design and deployment of autonomous AI agent systems that power Expedia's traveler experiences, working with large language models, tool integration, and multi-agent architectures. It suits experienced ML scientists who have shipped agentic systems to production and want to influence how millions of travelers plan and book trips.

Listed on Expedia Group’s careers site · Apply there ↗

Open-Source Machine Learning Engineer - US Remote

Hugging Face

Full-time · New York, NY · Remote

This role involves maintaining and improving Hugging Face's widely-used open-source machine learning libraries like Transformers and Datasets, collaborating directly with a global community of researchers and practitioners. It suits developers with a proven track record contributing to open-source ML projects who are passionate about making machine learning tools accessible and want to work on impactful ecosystem-level work.

Listed on Hugging Face’s careers site · Apply there ↗

View all 42 on Jobs

Asked for alongside JAX

Measured from the 42 open roles that name JAX — not from a curated list.

What employers state they pay

Median

$283,000

Middle half

$208,000 – $429,000

From 26 of 42 open roles — the 62% that publish a range, annualized to USD. Only some employers state pay, and the ones that do skew to California and New York, where the law requires it.

Who employers are hiring

  • Entry14%
  • Mid14%
  • Senior57%
  • Lead / Staff / Principal14%

From 28 roles whose posting states a level.

Where JAX roles are

Among the 28 of 42 open roles we could place on a map.

Employers hiring for JAX

16 in total, most open roles first.

Related skills

Curated neighbors in the taxonomy, whether or not employers ask for them together.

Turn on analytics and we load Google Analytics: Google gets the pages you open and what you do here — searches, jobs you view, jobs you apply to — and sets its own cookies. Leave it off and the only cookies we set are your login, your theme, and this answer. Accept All also records a yes to advertising, which nothing uses yet. Privacy Policy.