JAX Jobs - Remote & On-site

JAX jobs at vetted startups and product companies focused on high-performance ML and research, updated daily. These roles often pay competitively and support remote or hybrid work arrangements. Browse and apply today.

AI/ML
Data Science
Backend
Frontend
Mobile
DevOps
Web3
Games
Hardware
Robotics
Security
QA
Executive
Networking
Product
Design
Analytics
Support
Enterprise Apps
Quantum

Granica

granica.ai
Granica is an enterprise AI efficiency platform designed to reduce storage, compute, and context costs for large-scale AI applications. Founded in 2019 and headquartered in Mountain View, California, the company operates within its customers' own cloud perimeters (AWS and Google Cloud) to optimize enterprise data and model workflows securely without exposing sensitive information.
granica.ai • HQ: Mountain View, United States • Data & Analytics • Big Data • Artificial Intelligence • Est. 2019
HQ: Mountain View, United States • Data & Analytics • Big Data • Artificial Intelligence • Est. 2019
Verified live · 2 hours ago 2 months ago

Research Scientist – Diffusion Models

$145k – $318k per year (Estimated)In office • PhD • Full-Time
Python
AI/ML
AI Agents
Diffusion Models
JAX
Multimodal AI
PyTorch
Apply
Top 25% payVerified live · 2 hours ago 2 months ago

Research Scientist – Large Tabular Models (LTMs)

$160k – $250k per year • In office • PhD • Full-Time
Python
AI/ML
AI Agents
JAX
PyTorch
Apply
Report

Extropic

extropic.ai
Extropic is building thermodynamic computing hardware that performs probabilistic sampling in physics rather than in software. Founded in 2022 by former Google quantum computing researchers, it targets generative models that are natively stochastic. The company published its first prototype chips and an open simulation stack.
extropic.ai • HQ: New York, United States • 51-200 employees • Est. 2022
HQ: New York, United States • 51-200 employees • Est. 2022
Top 25% payVerified live · 2 days ago 3 months ago

Research Scientist - Machine Learning

$150k – $250k per year • In office (San Francisco, United States) • Full-Time
Python
AI/ML
Diffusion Models
JAX
Keras
ONNX
PyTorch
Ray
TensorFlow
Weights & Biases
DevOps
AWS
SLURM
Apply
Open 131 dayTop 25% pay 5 months ago

Research Scientist - Machine Learning

$150k – $250k per year • In office (Boston, United States) • Full-Time
Python
AI/ML
Diffusion Models
JAX
Keras
ONNX
PyTorch
Ray
TensorFlow
Weights & Biases
DevOps
AWS
SLURM
Apply
Report

Skild

skild.ai
Bringing AI into the physical world. Most AI today exists solely in cyberspace. Physical reality is far more information rich, and automating physical tasks offers immense returns. We are building AI that will interact with and affect change in the real world.
skild.ai • San Mateo • Bengaluru • Pittsburgh • San Francisco • Robotics • Artificial Intelligence • Robotics AI
San Mateo • Bengaluru • Pittsburgh • San Francisco • Robotics • Artificial Intelligence • Robotics AI
Open 263 daysVerified live · 1 day ago 9 months ago

Computer Vision AI & ML Engineer

$31k – $130k per year (Estimated)In office (Bengaluru, India)
C++
Python
C++
PyTorch C++
TensorFlow C++
AI/ML
Computer Vision
Multimodal AI
JAX
PyTorch
TensorFlow
Robotics
Gazebo
Isaac Sim
Sensor Fusion
Design
Blender
Apply
Open 707 daysVerified live · 1 day ago 2 years ago

Software Engineer, AI Training and Infrastructure

$95k – $229k per year (Estimated) • Middle • 3+ years expIn office (San Mateo, United States) • Bachelor's Degree • Middle
C++
Python
C++
PyTorch C++
TensorFlow C++
AI/ML
Edge AI
JAX
PyTorch
TensorFlow
DevOps
CI/CD
AWS
Azure
GCP
Apply
Open 707 daysVerified live · 1 day ago 2 years ago

Software Engineer, AI Inference

$95k – $229k per year (Estimated) • Middle • 3+ years expIn office (San Mateo, United States) • Bachelor's Degree • Middle
C++
Python
Rust
C++
PyTorch C++
TensorFlow C++
AI/ML
AI Agents
Edge AI
CUDA Toolkit
JAX
PyTorch
TensorFlow
CUDA
Apply
Report

Valsoft Corporation

valsoftcorp.com
Valsoft Corporation is a Canadian software holding company specializing in the acquisition, development, and management of vertical market software (VMS) businesses. Headquartered in Montreal, Quebec, the company operates under a decentralized, "hold-forever" philosophy, partnering with niche software firms to drive long-term growth and stability.
valsoftcorp.com • HQ: Montreal, Canada • Information Technology • Software • 1001-5000 employees
HQ: Montreal, Canada • Information Technology • Software • 1001-5000 employees
Verified live · 3 hours ago 3 months ago

Senior Java Developer - Chordline Health

$136k – $225k per year (Estimated) • Senior • Remote (United States) • Full-Time • Senior
JavaScript
SQL
Java
Node JS
Java
Apache Tomcat
Jackson
Node JS
Dotenv
Databases
MS SQL
AI/ML
JAX
Phi
Frontend
JQuery
Mobile
MVVM
DevOps
CI/CD
AWS
Azure
Docker
GCP
GitHub Actions
TeamCity
GitHub
Cybersecurity
HIPAA
QA
Playwright
Swagger
Apply
Report

SandboxAQ

sandboxaq.com
SandboxAQ leverages the compound effects of AI and advanced computing to address some of the biggest challenges impacting society. SandboxAQ technologies include AI simulation, cryptography management for cybersecurity, and AI sensing for global organizations.
sandboxaq.com • HQ: Palo Alto, United States • Cybersecurity • Quantum Science & Computing • Artificial Intelligence • Est. 2016
HQ: Palo Alto, United States • Cybersecurity • Quantum Science & Computing • Artificial Intelligence • Est. 2016
Top 25% payVerified live · 16 min ago 3 months ago

Senior Research Scientist, Battery Materials Simulation

$134k – $252k per year • Senior • 5+ years expRemote/Hybrid • PhD • Full-Time • Senior
Python
AI/ML
Braintrust
JAX
PyTorch
TensorFlow
Apply
Verified live · 16 min ago 4 months ago

Senior Machine Learning Engineer, AI Generation Engine

$86k – $152k per year • Senior • 5+ years exp • Remote (Canada) • Bachelor's Degree • Full-Time • Senior
Python
AI/ML
Braintrust
Fine-tuning
JAX
MLFlow
NumPy
Pandas
PyTorch
TensorFlow
Weights & Biases
Edge AI
Pre-training
AI Agents
DevOps
CI/CD
AWS
GCP
Apply
Verified live · 16 min ago 4 months ago

Staff Machine Learning Engineer, AI Generation Engine

$106k – $186k per year • Staff+ • 8+ years exp • Remote (Canada) • Bachelor's Degree • Full-Time • Staff
Python
AI/ML
Braintrust
Fine-tuning
JAX
MLFlow
NumPy
Pandas
PyTorch
TensorFlow
Weights & Biases
Edge AI
Pre-training
AI Agents
DevOps
CI/CD
AWS
GCP
Apply
Report

Pismo

pismo.io
Pismo provides cloud-native core banking and card processing infrastructure. Its platform runs accounts, cards and payments for banks worldwide. The company was acquired by a global payments network.
pismo.io • HQ: São Paulo, Brazil • FinTech • Financial Services • Banking • 1001-5000 employees • Est. 2016
HQ: São Paulo, Brazil • FinTech • Financial Services • Banking • 1001-5000 employees • Est. 2016
Verified live · 1 day ago 3 months ago

Senior Software Engineer - AI Agents

$123k – $191k per year • Senior • 5+ years expRemote/Hybrid (Ashburn, Foster City, United States) • Bachelor's Degree • Full-Time • Senior
Go
Java
Python
Rust
Databases
Cassandra
FAISS
MySQL
Pinecone
AI/ML
LLM Guardrails
AI Agents
AutoGen
Claude
CUDA
CUDA Toolkit
Embeddings
Fine-tuning
JAX
LangChain
LangGraph
Prompt Engineering
PyTorch
TensorFlow
DevOps
AWS
Azure
Docker
GCP
Kubernetes
QA
Postman
Apply
Report

Merlin Labs

merlinlabs.com
Merlin Labs is a company focused on developing autonomous flight capabilities through its innovative Merlin Pilot technology. They aim to deliver reduced aircrew and uncrewed flight solutions for various aircraft, including the C-130J Super Hercules, KC-135 Stratotanker, and F-16. Merlin Labs employs advanced automation systems, sensor technology, and natural language processing to enable full takeoff to touchdown flights, navigate environments, and communicate effectively with air traffic control.
merlinlabs.com • Boston • Washington
Boston • Washington
Top 25% payVerified live · 9 hours ago 3 months ago

Staff Software Engineer, AI Foundation Model

$225k – $390k per year • Staff+ • 8+ years expRemote/Hybrid (Boston, United States) • Full-Time • Staff
Python
C++
Rust
C++
PyTorch C++
TensorFlow C++
AI/ML
JAX
MLFlow
Weights & Biases
PyTorch
Quantization
Synthetic Data
TensorFlow
TensorRT
Transformers
Robotics
Sim-to-Real
Apply
Report

Amgen

amgen.com
Amgen is a global biopharmaceutical pioneer headquartered in Thousand Oaks, California, that specializes in discovering, developing, and manufacturing innovative biologic therapies. The company focuses on treating serious illnesses with high unmet medical needs across key areas including oncology, cardiovascular disease, inflammation, rare diseases, and nephrology. Leveraging advanced human genetics, molecular engineering, and biosimilar development, it serves millions of patients worldwide through established blockbuster treatments and cutting-edge pipelines.
amgen.com • HQ: Thousand Oaks, United States • Biotechnology • Health Care • Pharmaceuticals • 1001-5000 employees • Est. 1980
HQ: Thousand Oaks, United States • Biotechnology • Health Care • Pharmaceuticals • 1001-5000 employees • Est. 1980
Verified live · 1 day ago 3 months ago

Principal Machine Learning Engineer - Forecasting

$32k – $76k per year (Estimated) • Lead • 12+ years expIn office (Hyderabad, India) • Full-Time • Principal
Python
SQL
Python
FastAPI
AI/ML
AI Agents
Dagster
JAX
LLM
MLFlow
NLP
Prefect
PyTorch
Ray
Scikit-learn
Spark
TensorFlow
Human-in-the-Loop
LLM Guardrails
Function Calling
SGLang
TensorRT
TensorRT-LLM
vLLM
Model Context Protocol
RAG
DevOps
CI/CD
Analytics
A/B Testing
Apply
Verified live · 1 day ago 3 months ago

Senior Manager, Data Science - Forecasting

$29k – $67k per year (Estimated) • Lead • 15+ years expIn office (Hyderabad, India) • Full-Time • Senior
Python
SQL
Databases
Amazon Redshift
AI/ML
AI Agents
Dagster
JAX
LLM
MLFlow
Prefect
PyTorch
Scikit-learn
Spark
TensorFlow
Human-in-the-Loop
LLM Guardrails
Function Calling
Amazon SageMaker
Model Context Protocol
RAG
DevOps
Kubernetes
Amazon EC2
Amazon EKS
AWS
AWS Lambda
Azure
GCP
Amazon ECS
Amazon Kinesis
Amazon S3
Analytics
A/B Testing
Apply
Report

Engram

engram.ai
engram.ai • San Francisco
San Francisco
Verified live · 1 day ago 3 months ago

ML Systems & Performance Engineer

$163k – $316k per year (Estimated) • Senior • 5+ years expIn office (San Francisco, United States) • Full-Time • Senior
Edge AI
JAX
PyTorch
DevOps
Platform Engineering
Management
Notion
Apply
Report

Accellor

accellor.com
Accellor is an AI-first digital transformation and technology services firm headquartered in Frisco, Texas, and founded in 2018. The company develops purpose-built artificial intelligence solutions, including agentic AI for supply chain optimization, AI-assisted software engineering, and intelligent automation for sectors like healthcare and financial services. It operates globally with a focus on scaling AI from proof of concept to enterprise-wide production through strategic partnerships with major technology providers like Salesforce and Microsoft.
accellor.com • HQ: Frisco, United States • Health Care • Business Process Automation (BPA) • Software • AI Agents • LLM & Generative AI • Artificial Intelligence • Information Technology • IT Consulting • Est. 2018
HQ: Frisco, United States • Health Care • Business Process Automation (BPA) • Software • AI Agents • LLM & Generative AI • Artificial Intelligence • Information Technology • IT Consulting • Est. 2018
Verified live · 1 day ago 3 months ago

AI Principal Engineer

$162k – $346k per year (Estimated) • Lead • 10+ years expIn office (San Francisco, United States) • Principal
C++
Go
Java
Python
Rust
TypeScript
C++
PyTorch C++
TensorFlow C++
AI/ML
AI Agents
ChatGPT
CUDA Toolkit
Function Calling
JAX
LLM
Multimodal AI
PyTorch
Ray
Reinforcement Learning
TensorFlow
vLLM
CUDA
Triton
Context Engineering
LLM Evaluation
NCCL
OpenAI
OpenAI Codex
Post-training
Pre-training
Hallucination
RAG
DevOps
Kubernetes
Platform Engineering
CI/CD
Docker
Grafana
OpenTelemetry
Prometheus
Terraform
Apply
Report

MacPaw

macpaw.com
MacPaw is a software company headquartered in Kyiv, Ukraine, and founded in 2008 by Oleksandr Kosovan. The company develops consumer applications for macOS and iOS including CleanMyMac, the Setapp subscription app service, Encrypto, and ClearVPN, sold directly to users worldwide. It employs several hundred people, has continued operating through the war in Ukraine, and runs the MacPaw Foundation alongside its commercial work.
macpaw.com • HQ: Kyiv, Ukraine • Operating Systems • Endpoint Security • Laptops • Consumer Electronics • Productivity Software • Software • Cybersecurity • Est. 2008
HQ: Kyiv, Ukraine • Operating Systems • Endpoint Security • Laptops • Consumer Electronics • Productivity Software • Software • Cybersecurity • Est. 2008
Verified live · 1 day ago 3 months ago

AI Research Scientist

Remote (Ukraine) • Master's Degree
C++
Python
C++
PyTorch C++
TensorFlow C++
AI/ML
Fine-tuning
JAX
Knowledge Distillation
LLM
LoRA
MLX ML
NLP
PyTorch
Quantization
TensorFlow
PEFT
Edge AI
Mixture of Experts
Management
Jira
Apply
Report

SatSure

satsure.co
SatSure builds decision analytics from satellite imagery and other geospatial data. Its products serve agricultural lending, insurance, infrastructure and climate risk. Banks and government bodies use its crop and land intelligence.
satsure.co • HQ: Bengaluru, India • Data & Analytics • Decision Intelligence • Artificial Intelligence • Est. 2017
HQ: Bengaluru, India • Data & Analytics • Decision Intelligence • Artificial Intelligence • Est. 2017
$28k – $77k per year (Estimated) • Middle • 3+ years expIn office (Bengaluru, India) • Full-Time • Middle
Contrastive Learning
DeepSpeed
Fine-tuning
Mamba
Multimodal AI
Segment Anything
Self-Supervised Learning
FSDP
Time Series Forecasting
CUDA
CUDA Toolkit
JAX
Mixture of Experts
PyTorch
Apply
Open 108 days 4 months ago

Senior ML Researcher – Foundation Models

$33k – $83k per year (Estimated) • Senior • 5+ years expIn office (Bengaluru, India) • Full-Time • Senior
Contrastive Learning
DeepSpeed
Fine-tuning
Mamba
Multimodal AI
Segment Anything
Self-Supervised Learning
FSDP
Time Series Forecasting
CUDA
CUDA Toolkit
JAX
Mixture of Experts
PyTorch
Apply
Report

Writer

writer.com
Writer is an enterprise generative artificial intelligence company headquartered in San Francisco, California, and founded in 2020. The company provides a full-stack platform that includes proprietary large language models, an integrated graph-based retrieval-augmented generation system, and tools for building autonomous AI agents. It serves Fortune 500 companies across sectors such as financial services, healthcare, and retail, focusing on delivering secure and compliant AI applications for marketing, sales, and operations.
writer.com • HQ: San Francisco, United States • Business Process Automation (BPA) • Software • Marketing Automation • Sales & Marketing • LLM & Generative AI • AI Agents • Artificial Intelligence • Est. 2020
HQ: San Francisco, United States • Business Process Automation (BPA) • Software • Marketing Automation • Sales & Marketing • LLM & Generative AI • AI Agents • Artificial Intelligence • Est. 2020
Top 25% payVerified live · 9 hours ago 3 months ago

AI research scientist

$234k – $349k per year • Equity • Senior • 7+ years expRemote/Hybrid (San Francisco, New York, Seattle, United States) • PhD • Full-Time • Senior
Python
AI/ML
AI Agents
Fine-tuning
JAX
LLM
NLP
PyTorch
Reinforcement Learning
RLHF
Synthetic Data
DPO
GRPO
Post-training
Pre-training
SFT
Function Calling
Apply
Top 25% payVerified live · 9 hours ago 4 months ago

AI engineer

$168k – $329k per year • Equity • Senior • 5+ years expRemote/Hybrid (San Francisco, New York, United States) • Full-Time • Senior
Python
AI/ML
AI Agents
JAX
PyTorch
TensorFlow
DevOps
AWS
Azure
GCP
Apply
Verified live · 9 hours ago 4 months ago

AI engineer (UK)

$107k – $221k per year (Estimated) • Equity • Senior • 5+ years expRemote/Hybrid (London, United Kingdom) • Full-Time • Senior
Python
AI/ML
AI Agents
JAX
PyTorch
TensorFlow
DevOps
AWS
Azure
GCP
Apply
Report

Menlo Research

menlo.ai
Menlo Research is an applied AI research lab and robotics software startup headquartered in Singapore with operations globally. Founded in 2023, the company builds open-source AI infrastructure, modular operating software (Asimov OS), and world modeling systems that function as the "brain" for next-generation humanoid and mobile robots. Through an open-innovation approach and shared API architecture, Menlo enables developers and OEMs to train, deploy, and monetize composable robotic skills and multimodal perception models across diverse hardware platforms.
menlo.ai • HQ: San Francisco, United States
HQ: San Francisco, United States
Verified live · 2 hours ago 3 months ago

Robotics AI Engineer

In office (Ho Chi Minh City, Vietnam) • Full-Time
Python
AI/ML
JAX
PyTorch
Reinforcement Learning
Robotics
Imitation Learning
Isaac Gym
Isaac Lab
MuJoCo
Reinforcement Learning
Sim-to-Real
Isaac Sim
Apply
Report

1X Technologies

1x.tech
1X Technologies is a humanoid robotics company founded in Norway in 2014 that develops general-purpose robots for everyday physical work. Its EVE wheeled android and the newer NEO home humanoid are trained with end-to-end neural networks rather than scripted motion, and the company runs its own hardware manufacturing. Backed by OpenAI and other investors, it operates engineering teams in Norway and the San Francisco Bay Area.
1x.tech • HQ: Moss, Norway • Robotics Hardware • Robotics AI • Robotics • Est. 2014
HQ: Moss, Norway • Robotics Hardware • Robotics AI • Robotics • Est. 2014
Top 25% payVerified live · 1 hour ago 3 months ago

Software Engineer - Simulation

$200k – $280k per year • Middle • 4+ years expIn office • Master's Degree • Full-Time • Middle
C++
Python
C++
OpenGL
AI/ML
Synthetic Data
JAX
Robotics
Isaac Sim
MuJoCo
PyBullet
Sim-to-Real
Apply
Report

Snowflake

snowflake.com
Snowflake is a cloud software company that helps organisations store, process, and analyse large volumes of data in a way that is designed for modern cloud infrastructure. Many businesses struggle with data spread across different systems, slow reporting cycles, and the cost and complexity of managing traditional data warehouses.
snowflake.com • HQ: Bozeman, United States • Database Management • Data & Analytics • Data Warehousing • 1001-5000 employees
HQ: Bozeman, United States • Database Management • Data & Analytics • Data Warehousing • 1001-5000 employees
Top 25% payVerified live · 2 hours ago 3 months ago

AI Research Scientist, New Grad – Agents & Reinforcement Learning

$176k – $230k per year • Junior • In office (Bellevue, United States) • PhD • Full-Time • Junior
Python
Databases
Snowflake
AI/ML
Fine-tuning
JAX
LLM
PyTorch
Reinforcement Learning
RLHF
DPO
Post-training
PPO
AI Agents
Apply
Report
Genesis Molecular AI is a drug discovery company headquartered in Burlingame, California, and founded in 2019. The company combines deep learning models of molecular structure and binding with in-house chemistry and biology to design small molecule drug candidates against difficult targets. It advances its own pipeline while running partnered discovery programmes with large pharmaceutical companies, and previously operated under the name Genesis Therapeutics.
genesis.ml • HQ: Burlingame, United States • Pharmaceuticals • Bioinformatics • Biopharma • Health Care • Drug Discovery • Medical AI • Artificial Intelligence • Biotechnology • Est. 2019
HQ: Burlingame, United States • Pharmaceuticals • Bioinformatics • Biopharma • Health Care • Drug Discovery • Medical AI • Artificial Intelligence • Biotechnology • Est. 2019
$99k – $257k per year (Estimated) • Equity • In office (San Mateo, United States) • PhD • Full-Time
Python
AI/ML
Diffusion Models
JAX
NumPy
PyTorch
Mixture of Experts
DevOps
HPC
Apply
Report

Goodfire

goodfire.com
Goodfire is an AI interpretability research company focused on understanding and designing AI systems you can trust. Our mission is to advance humanity's understanding of AI to build safe and powerful AI systems.
goodfire.com • London • San Francisco • Machine Learning • Biotechnology • Artificial Intelligence
London • San Francisco • Machine Learning • Biotechnology • Artificial Intelligence
Open 228 daysTop 25% pay 8 months ago

Head of Engineering, Platform

$300k – $450k per year • Executive • 10+ years expIn office (San Francisco, United States) • Architect
JAX
PyTorch
OpenAI
Apply
Open 139 daysTop 25% pay 5 months ago

Field Team - Member of Technical Staff

$200k – $325k per year • Staff+ • In office (San Francisco, United States) • Staff
Python
AI/ML
JAX
PyTorch
OpenAI
Fine-tuning
Apply
Report

Mach9

mach9.com
Mach 9 is a Digital marketing agency in NYC focusing on brands and industries. It's time to break free from your digital echo chambers.
mach9.com • San Francisco • Sales & Marketing • Advertising
San Francisco • Sales & Marketing • Advertising
Top 25% payVerified live · 1 day ago 3 months ago

ML Engineer

$180k – $300k per year • In office (San Francisco, United States) • Master's Degree • Full-Time
Python
AI/ML
Computer Vision
JAX
Knowledge Distillation
ONNX
PyTorch
Quantization
TensorFlow
TensorRT
Weights & Biases
Apply
Report

Blueorigin

blueorigin.com
Blue Origin is an American aerospace manufacturer, defense contractor, and spaceflight services company founded in 2000 by Jeff Bezos. Headquartered in Kent, Washington, the company designs and operates reusable rockets, heavy-lift engines, and deep-space logistics infrastructure to enable sustainable space exploration and commercial operations.
blueorigin.com • Los Angeles • Washington • Seattle • Huntsville • Arlington • Space Services • Space & Aerospace • Launch Services • 1001-5000 employees
Los Angeles • Washington • Seattle • Huntsville • Arlington • Space Services • Space & Aerospace • Launch Services • 1001-5000 employees
Open 92 daysTop 25% pay 3 months ago

Chief Engineer - ML, Physical AI, & Autonomy

$254k – $355k per year • Equity • Executive • 10+ years expIn office (Seattle, United States) • Master's Degree • Full-Time • Architect
C++
Python
MATLAB
C++
PyTorch C++
TensorFlow C++
MATLAB
Simulink
AI/ML
Computer Vision
JAX
PyTorch
TensorFlow
AI Agents
Robotics
Digital Twin
Sensor Fusion
Gazebo
Isaac Sim
MuJoCo
ROS
ROS2
Apply
Report
JAX Jobs - Remote & On-site
Frequently asked questions
JAX: How many jobs are available now?
There are 491 JAX jobs listed on Alion right now.
JAX: What is the average salary for these roles?
The average salary for JAX roles on Alion is $285,074 USD.
JAX: Are remote, relocation, or visa sponsorship options available?
Many JAX listings allow remote or hybrid schedules, and some positions may offer relocation or visa sponsorship-verify on each job page.
JAX: How do I apply for JAX jobs on Alion?
Open the JAX job posting and click Apply or follow the employer’s link; sign in to save jobs and enable alerts for new listings.