# JAX

Skill · Data Science, Analytics and AI/ML

Canonical page: https://career.thegoodapps.co/skills/jax

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

- Open roles: 50
- Employers hiring: 20
- Median stated pay: $240,000 (middle half $208,000 – $296,000), from 34 of 50 open roles — the 68% 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.
- Last checked: September 13, 2026

## Asked for alongside JAX

Measured from the 50 open roles that name JAX, not from a curated list.

- [PyTorch](https://career.thegoodapps.co/skills/pytorch) — 46 roles (92.0%)
- [Python](https://career.thegoodapps.co/skills/python) — 34 roles (68.0%)
- [TensorFlow](https://career.thegoodapps.co/skills/tensorflow) — 24 roles (48.0%)
- [CUDA](https://career.thegoodapps.co/skills/cuda) — 21 roles (42.0%)
- [C++](https://career.thegoodapps.co/skills/c-2) — 13 roles (26.0%)
- [Kubernetes](https://career.thegoodapps.co/skills/kubernetes) — 12 roles (24.0%)
- [Machine Learning](https://career.thegoodapps.co/skills/machine-learning) — 7 roles (14.0%)
- [Rust](https://career.thegoodapps.co/skills/rust) — 7 roles (14.0%)

## Seniority of these roles

- Entry: 4 (10%)
- Mid: 6 (15%)
- Senior: 23 (59%)
- Lead / Staff / Principal: 6 (15%)

From 39 roles whose posting states a level.

## Where JAX roles are

- [New York, NY](https://career.thegoodapps.co/jobs/in/new-york-ny) — 13
- [San Francisco, CA](https://career.thegoodapps.co/jobs/in/san-francisco-ca) — 8
- [Santa Clara, CA](https://career.thegoodapps.co/jobs/in/santa-clara-ca) — 2
- [Boston, MA](https://career.thegoodapps.co/jobs/in/boston-ma) — 1
- [Redwood City, CA](https://career.thegoodapps.co/jobs/in/redwood-city-ca) — 1
- [San Jose, CA](https://career.thegoodapps.co/jobs/in/san-jose-ca) — 1

Among the 27 of 50 open roles we could place on a map.

## Roles that use JAX

- [Software Engineer](https://career.thegoodapps.co/occupations/software-engineer) — 17
- [Machine Learning Engineer](https://career.thegoodapps.co/occupations/machine-learning-engineer) — 11
- [Data Scientist](https://career.thegoodapps.co/occupations/data-scientist) — 2
- [Engineering Manager](https://career.thegoodapps.co/occupations/engineering-manager) — 2
- [Intern](https://career.thegoodapps.co/occupations/intern) — 1
- [Product Marketing Manager](https://career.thegoodapps.co/occupations/product-marketing-manager) — 1

## Employers hiring for JAX

- [Cohere](https://career.thegoodapps.co/organizations/cohere) — 13 open roles
- [NVIDIA](https://career.thegoodapps.co/organizations/nvidia) — 11 open roles
- [Menlo Ventures Portfolio](https://career.thegoodapps.co/organizations/menlo-ventures-portfolio) — 5 open roles
- [Crusoe](https://career.thegoodapps.co/organizations/crusoe) — 2 open roles
- [Runway](https://career.thegoodapps.co/organizations/runway) — 2 open roles
- [Target](https://career.thegoodapps.co/organizations/target) — 2 open roles

Related skills: [PyTorch](https://career.thegoodapps.co/skills/pytorch), [TensorFlow](https://career.thegoodapps.co/skills/tensorflow), [NumPy](https://career.thegoodapps.co/skills/numpy)

## Open roles requiring JAX (50)

- [Senior Software Engineer, GPU Infrastructure (HPC)](https://career.thegoodapps.co/jobs/51d9d309-29a1-4f12-94b1-c43dd27daf6e/senior-software-engineer-gpu-infrastructure-hpc-at-cohere) at [Cohere](https://career.thegoodapps.co/organizations/cohere) — Full-time, Canada, United States · Remote, CA$285,000 – CA$340,000
- [Senior Software Engineer - AI Agents](https://career.thegoodapps.co/jobs/9ff1f8da-3e36-4c38-aa2b-c375ad1eb5d3/senior-software-engineer-ai-agents-at-visa-inc) at [Visa Inc.](https://career.thegoodapps.co/organizations/visa-inc) — Full-time, $123,000 – $190,900
- [Sr Applied Data Scientist - Personalization (applied ML, PyTorch, RecSys)](https://career.thegoodapps.co/jobs/15d1f6c6-6d30-4484-87df-95fd15ed8a34/sr-applied-data-scientist-personalization-applied-ml-pytorch-recsys-at-target) at [Target](https://career.thegoodapps.co/organizations/target) — 7000 Target Pkwy N,NCD-0375 Brooklyn Park,MN 55445, $98,000 – $176,000
- [Principal AI Researcher](https://career.thegoodapps.co/jobs/a4b54056-34d1-4ea3-934f-9df9a6d99f2e/principal-ai-researcher-at-workday) at [Workday](https://career.thegoodapps.co/organizations/workday) — Full-time, $228,000 – $342,000
- [Senior Software Engineer, CUDA C++ Core Libraries](https://career.thegoodapps.co/jobs/d1a26576-7b0e-48a9-bbf6-ef295140ff9f/senior-software-engineer-cuda-c-core-libraries-at-nvidia) at [NVIDIA](https://career.thegoodapps.co/organizations/nvidia) — $184,000 – $287,500
- [Senior AI Software Architect - Runtime](https://career.thegoodapps.co/jobs/f9a6eed3-de6d-4efd-9c49-b37ed3a91e99/senior-ai-software-architect-runtime-at-intel) at [Intel](https://career.thegoodapps.co/organizations/intel) — Full-time, $195,200 – $361,200
- [AI Robotics Researcher Intern (Dexterous Manipulation)](https://career.thegoodapps.co/jobs/08133ed8-6a04-4457-9f3f-2b5631cdbd26/ai-robotics-researcher-intern-dexterous-manipulation-at-nio) at [NIO](https://career.thegoodapps.co/organizations/nio) — Internship, San Jose, CA, $38 – $46 / hour
- [Staff ML Engineer, Fine Tuning - Slack](https://career.thegoodapps.co/jobs/be597aae-214c-475e-9854-ed37a11266c0/staff-ml-engineer-fine-tuning-slack-at-slack) at [Slack](https://career.thegoodapps.co/organizations/slack) — $197,300 – $313,700
- [Senior Machine Learning Scientist - Agentic Experience](https://career.thegoodapps.co/jobs/ed919295-f53b-4fb9-bc0d-732067dec813/senior-machine-learning-scientist-agentic-experience-at-expedia-group) at [Expedia Group](https://career.thegoodapps.co/organizations/expedia-group) — $173,000 – $242,500
- [Principal Machine Learning Engineer, Data Mining](https://motional.com/open-positions/?gh_jid=7779985003#/7779985003) at [Motional](https://career.thegoodapps.co/organizations/motional) — Full-time, Boston, MA · Remote, $144,000 – $192,000 (listed on Motional's own careers site)

All 50 open roles: https://career.thegoodapps.co/jobs?skills=16733cae-1fdc-4410-bea1-2dc6325e9e94