Software Engineer, JAX External

1 Month 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

ByteDance - Backend Software Engineer

ByteDance

Dubai, Dubai, United Arab Emirates (On-Site)
4 Months ago
ION - Software Developer/Engineer - Graduate Development Program

ION

Turin, Piedmont, Italy (On-Site)
7 Months ago
Microsoft - Software Engineer

Microsoft

(On-Site)
3 Weeks ago
Riot Games - Manager, Software Engineering

Riot Games

Shanghai, Shanghai, China (On-Site)
2 Months ago
Microsoft - Senior DPU Software Engineer - Secure Enclave

Microsoft

Bengaluru, Karnataka, India (On-Site)
1 Month ago
Google - Software Engineer II, Filestore Control Plane

Google

Tel Aviv-Yafo, Tel Aviv District, Israel (On-Site)
1 Month ago
Google - Software Engineer II, Device Integrity, Web and IOS

Google

Zürich, Zurich, Switzerland (On-Site)
1 Month ago
Kaedim - Customer Support Engineer

Kaedim

San Francisco, California, United States (On-Site)
9 Months ago
Google - Software Engineer, Early Career, Chrome

Google

Mexico City, Mexico City, Mexico (On-Site)
1 Month ago
Google - Software Engineer III, CorpEng

Google

Hyderabad, Telangana, India (On-Site)
1 Month ago

Get notifed when new similar jobs are uploaded

Similar Skill Jobs

WebFX - Sr  Quality Assurance Engineer (Philippines )

WebFX

Philippines (Remote)
6 Months ago
Google - Software Engineering Manager, Google Play Games

Google

Bengaluru, Karnataka, India (On-Site)
1 Month ago
Google - Senior Software Engineer, Global Network Edge

Google

Dublin, County Dublin, Ireland (On-Site)
3 Weeks ago
Velotio Technologies - Senior GenAI Engineer - Java

Velotio Technologies

Maharashtra, India (Remote)
1 Month ago
ByteDance - Senior Software Engineer, Global Payment Risk & Compliance

ByteDance

San Jose, California, United States (On-Site)
6 Months ago
Google - Application Engineer, SAP Specialist

Google

Bengaluru, Karnataka, India (On-Site)
1 Month ago
Hive Innovative Group - Frontend developer

Hive Innovative Group

Cairo, Cairo Governorate, Egypt (On-Site)
10 Months ago
Rackspace Technology - R-19462 Data Engineer III - VN

Rackspace Technology

Vietnam (Remote)
4 Months ago
The Walt Disney Company - Lead Software Engineer - iOS/Android

The Walt Disney Company

Seattle, Washington, United States (On-Site)
1 Month ago
SmileGate - Platform Division Web Service Backend Developer

SmileGate

Seongnam-si, Gyeonggi-do, South Korea (On-Site)
4 Months ago

Get notifed when new similar jobs are uploaded

Jobs in San Francisco, California, United States

On Location - Senior Accountant - Olympics

On Location

Atlanta, Georgia, United States (On-Site)
4 Months ago
ByteDance - Software Engineer Graduate (XR Web Platform-PICO)

ByteDance

San Jose, California, United States (On-Site)
1 Month ago
Brightline - Unity Developer

Brightline

Virginia, United States (Hybrid)
1 Month ago
Sports radar - Synergy Sports Showcase Softball/Baseball Video Scouts

Sports radar

Alabama, New York, United States (On-Site)
6 Months ago
Google - Data Scientist III, Product, Google Wallet

Google

Mountain View, California, United States (On-Site)
1 Month ago
ION - Senior Linux Systems Administrator - Somerset, NJ

ION

Clifton, New Jersey, United States (Hybrid)
7 Months ago
Nintendo - Principal, Partner Management - Second Party

Nintendo

Redmond, Washington, United States (Hybrid)
5 Months ago
Next Level Business Services - UX/UI Developer

Next Level Business Services

Cincinnati, Ohio, United States (On-Site)
7 Months ago
AI Fund - Chief Executive Officer

AI Fund

United States (Remote)
1 Month ago
Nintendo - Sr Bilingual Communications Coordinator - Japanese

Nintendo

Redmond, Washington, United States (Hybrid)
5 Months ago

Get notifed when new similar jobs are uploaded

Full Stack Development Jobs

The Walt Disney Company - Lead Software Engineer

The Walt Disney Company

Connecticut, United States (On-Site)
3 Months ago
Nagarro - Senior Staff Engineer

Nagarro

Mexico (Remote)
7 Months ago
Google - Software Engineer, Machine Learning

Google

New Taipei, New Taipei City, Taiwan (On-Site)
1 Month ago
Tamatem Games - Software Engineer

Tamatem Games

Amman, Amman Governorate, Jordan (On-Site)
1 Month ago
GT - Senior Full-stack Developer (Python+React.js)

GT

Ukraine (Remote)
1 Month ago
Shaadi - Senior Software Engineer Node.js

Shaadi

Maharashtra, India (On-Site)
8 Months ago
ByteDance - Service Framework Software Engineer Intern

ByteDance

Singapore (On-Site)
1 Month ago
Google - Software Engineer III, Performance, Platforms Infrastructure Engineering

Google

Sunnyvale, California, United States (On-Site)
1 Month ago
Playrix - Senior C++ Software Engineer (Build System)

Playrix

Serbia (Remote)
6 Months ago
Team Liquid - Senior Back End Engineer, EU

Team Liquid

Utrecht, Utrecht, Netherlands (Hybrid)
2 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.

London, England, United Kingdom (On-Site)

Bengaluru, Karnataka, India (On-Site)

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

Bengaluru, Karnataka, India (On-Site)

New York, New York, United States (On-Site)

Waterloo, Ontario, Canada (On-Site)

Taipei City, Taiwan (On-Site)

Austin, Texas, 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