Software Engineer, JAX External

6 Days ago • 2-4 Years • Full Stack Development • $141,000 PA - $202,000 PA

Job Summary

Job Description

This role involves building open-source infrastructure for cutting-edge ML on a small, high-agency team. You'll collaborate with engineers across JAX, XLA, and other teams, contributing to impactful projects while advocating open-source software and engaging with the ML community. Responsibilities include adding new features and publishing high-performance open-source Pallas kernels; implementing and benchmarking reference JAX models and techniques; working with the JAX team to design and implement new features; and directly assisting JAX power users to solve complex problems and enable new workloads. The ideal candidate will have experience in software development, testing, and design, along with expertise in areas like networking, systems architecture, or AI infrastructure. Strong programming skills in languages like C, C++, Java, or Python are essential.
Must have:
  • Bachelor's degree or equivalent experience
  • 2+ years software development experience
  • Experience with data structures/algorithms
  • Experience testing, maintaining, or launching software
  • Software design and architecture experience
Good to have:
  • 4+ years experience in networking, storage, etc.
  • Experience with C, C++, Java, or Python
  • Building high-quality AI infrastructure
  • Designing clean and composable APIs
  • Experience with low-level C++ and high-level Python
  • Knowledge of latest AI models and techniques

Job Details


Minimum qualifications:

  • Bachelor’s degree or equivalent practical experience.
  • 2 years of experience with software development in one or more programming languages, and with data structures/algorithms.
  • 2 years of experience testing, maintaining, or launching software products, and 1 year of experience with software design and architecture.

Preferred qualifications:

  • 4 years of experience in networking, storage, systems architecture, compilers, operating systems, or modeling and analysis.
  • Experience programming in C, C++, Java, or Python.
  • Experience building high-quality and reusable infrastructure for AI.
  • Experience designing clean and composable APIs.
  • Experience diving into stack-spanning systems and tools, from high-level Python to low-level C++.
  • Ability to stay up-to-date on the latest AI models, tools, and techniques.

About the job

Google's software engineers develop the next-generation technologies that change how billions of users connect, explore, and interact with information and one another. Our products need to handle information at massive scale, and extend well beyond web search. We're looking for engineers who bring fresh ideas from all areas, including information retrieval, distributed computing, large-scale system design, networking and data storage, security, artificial intelligence, natural language processing, UI design and mobile; the list goes on and is growing every day. As a software engineer, you will work on a specific project critical to Google’s needs with opportunities to switch teams and projects as you and our fast-paced business grow and evolve. We need our engineers to be versatile, display leadership qualities and be enthusiastic to take on new problems across the full-stack as we continue to push technology forward.

With your technical expertise you will manage project priorities, deadlines, and deliverables. You will design, develop, test, deploy, maintain, and enhance software solutions.

In this role, you will build open-source infrastructure for cutting-edge ML on a small, high-agency team. You'll collaborate closely with engineers across the JAX, XLA, and other adjacent teams, and have the freedom to pursue impactful projects while advocating open-source software and engaging directly with the ML community.

Google Cloud accelerates every organization’s ability to digitally transform its business and industry. We deliver enterprise-grade solutions that leverage Google’s cutting-edge technology, and tools that help developers build more sustainably. Customers in more than 200 countries and territories turn to Google Cloud as their trusted partner to enable growth and solve their most critical business problems.

The US base salary range for this full-time position is $141,000-$202,000 + bonus + equity + benefits. Our salary ranges are determined by role, level, and location. Within the range, individual pay is determined by work location and additional factors, including job-related skills, experience, and relevant education or training. Your recruiter can share more about the specific salary range for your preferred location during the hiring process.

Please note that the compensation details listed in US role postings reflect the base salary only, and do not include bonus, equity, or benefits. Learn more about .

Responsibilities

  • Add new features and publish high-performance open-source Pallas kernels.
  • Implement and benchmark reference JAX models and techniques, and inform new JAX features and improvements.
  • Work with the JAX team to design and implement new JAX features.
  • Work directly with JAX power users to solve tricky problems and enable new workloads.

Similar Jobs

Knuddels - Software Development Internship

Knuddels

Baden-Württemberg, Germany (On-Site)
4 Weeks ago
Google - Video Streaming Specialist

Google

Hyderabad, Telangana, India (On-Site)
1 Week ago
Netflix - Data Engineer 5 - Playback

Netflix

United States (Remote)
1 Week ago
Google - Software Engineering Manager II, Google Cloud Kubernetes Networking

Google

Sunnyvale, California, United States (On-Site)
6 Days ago
ByteDance - Site Reliability Engineer - CapCut - San Jose/Seattle

ByteDance

San Jose, California, United States (On-Site)
5 Months ago
Nagarro - Staff Engineer, Java

Nagarro

Riyadh, Riyadh Province, Saudi Arabia (On-Site)
5 Months ago
The Walt Disney Company - Software Engineer II - C++

The Walt Disney Company

San Francisco, California, United States (On-Site)
1 Week ago
Forgotten Empires - Senior Software Engineer - Age of Empires (Remote)

Forgotten Empires

Prague, Czechia (Remote)
4 Weeks ago
Next Level Business Services - .NET Developer

Next Level Business Services

Chicago, Illinois, United States (On-Site)
6 Months ago

Get notifed when new similar jobs are uploaded

Similar Skill Jobs

Warner Bros Games - Senior Data Engineer

Warner Bros Games

Atlanta, Georgia, United States (Hybrid)
2 Months ago
Garena - Quality Assurance Intern

Garena

Singapore (On-Site)
1 Week ago
PwC - Data Architect – Technology Consulting

PwC

Prague, Prague, Czechia (On-Site)
6 Months ago
DraftKings - Lead Site Reliability Engineer

DraftKings

Boston, Massachusetts, United States (On-Site)
1 Month ago
Google - Technical Solutions Engineer, Infrastructure, GKE/Anthos

Google

Maharashtra, India (On-Site)
1 Week ago
Nordcurrent - Senior Gameplay Programmer

Nordcurrent

Vilnius, Vilnius County, Lithuania (On-Site)
5 Months ago
Epic Games - Senior Gameplay Systems Programmer, Fortnite

Epic Games

London, England, United Kingdom (On-Site)
2 Months ago
Next Level Business Services - Full Stack Developer

Next Level Business Services

Jersey City, New Jersey, United States (On-Site)
6 Months ago
The Walt Disney Company - Lead Software Engineer, Ad Platforms

The Walt Disney Company

San Francisco, California, United States (On-Site)
1 Month ago
Rackspace Technology - Oracle Application Developer II

Rackspace Technology

Gurugram, Haryana, India (Remote)
1 Week ago

Get notifed when new similar jobs are uploaded

Jobs in San Francisco, California, United States

Google - Software Engineer III, Core Machine Learning, Google Cloud

Google

Mountain View, California, United States (On-Site)
5 Months ago
Aristocrat Gaming - Senior Systems Reliability Engineer (SRE)

Aristocrat Gaming

Austin, Texas, United States (Hybrid)
1 Month ago
Tencent - Video Standards & Encoding Optimization Intern

Tencent

Palo Alto, California, United States (On-Site)
5 Months ago
The Walt Disney Company - Senior Software Engineer (Roku Engineer)

The Walt Disney Company

Santa Monica, California, United States (On-Site)
3 Months ago
The Walt Disney Company - Manager, International Disney+ Subscriber Planning

The Walt Disney Company

Santa Monica, California, United States (On-Site)
4 Months ago
Next Level Business Services - Sr. Python Admin

Next Level Business Services

Saint Paul, Minnesota, United States (On-Site)
6 Months ago
Fliff  Inc  - Operations Intern

Fliff Inc

Austin, Texas, United States (On-Site)
9 Months ago
Next Level Business Services - MDM IDQ Developer/ Solution Architect - Full Time Only

Next Level Business Services

New York, New York, United States (On-Site)
6 Months ago
The Walt Disney Company - Content Programming Associate Manager

The Walt Disney Company

Glendale, California, United States (Hybrid)
4 Days ago
Dun & Bradstreet - Account Executive, Outbound - Manage (R-16769)

Dun & Bradstreet

Jacksonville, Florida, United States (On-Site)
6 Months ago

Get notifed when new similar jobs are uploaded

Full Stack Development Jobs

The Walt Disney Company - Lead IT Developer

The Walt Disney Company

Montévrain, Île-de-France, France (On-Site)
4 Days ago
Netflix - Full-Stack Engineer (L5)

Netflix

Warsaw, Masovian Voivodeship, Poland (On-Site)
1 Week ago
Hasura - Software Engineer, Data Connectors (Bangalore, India)

Hasura

Bengaluru, Karnataka, India (Hybrid)
6 Months ago
Nagarro - Engineer, InfraOps

Nagarro

India (Remote)
6 Months ago
PhonePe - Software Engineer (Backend, 5-8 Years) - Financial Services

PhonePe

Bengaluru, Karnataka, India (On-Site)
5 Months ago
Google - Web Solutions Engineer I, Tools Development and Automation

Google

Hyderabad, Telangana, India (On-Site)
1 Week ago
Milestone - Senior Fullstack Developer

Milestone

Copenhagen, Denmark (On-Site)
1 Month ago
Scopely - Senior Software Engineer

Scopely

Spain (Hybrid)
5 Months ago
N-iX - Middle .NET Fullstack Engineer

N-iX

Colombia (Remote)
1 Week ago
Meta - Software Engineer, Infrastructure

Meta

Menlo Park, California, United States (Remote)
5 Months ago

Get notifed when new similar jobs are uploaded

About The Company

A problem isn't truly solved until it's solved for all. Googlers build products that help create opportunities for everyone, whether down the street or across the globe. Bring your insight, imagination and a healthy disregard for the impossible. Bring everything that makes you unique. Together, we can build for everyone.

Fremont, California, United States (On-Site)

Mountain View, California, United States (On-Site)

Bengaluru, Karnataka, India (On-Site)

Dublin, County Dublin, Ireland (On-Site)

Atlanta, Georgia, United States (On-Site)

Ho Chi Minh City, Ho Chi Minh City, Vietnam (On-Site)

Seattle, Washington, United States (On-Site)

View All Jobs

Get notified when new jobs are added by Google

Level Up Your Career in Game Development!

Transform Your Passion into Profession with Our Comprehensive Courses for Aspiring Game Developers.

Job Common Plug