Software Engineer, JAX External

1 Month ago • 2-4 Years • Full Stack Development • $141,000 PA - $202,000 PA

Job Summary

Job Description

Google is seeking a Software Engineer to build open-source infrastructure for cutting-edge ML on a small, high-agency team. Collaborate with engineers across JAX, XLA, and other teams. Responsibilities include adding new features and publishing high-performance open-source Pallas kernels, implementing and benchmarking reference JAX models, designing and implementing new JAX features, and working directly with JAX power users to solve problems and enable new workloads. The role requires strong software development skills, experience with data structures/algorithms, testing, maintaining, or launching software products, and software design and architecture. Experience with networking, storage, systems architecture, compilers, or AI infrastructure is highly preferred. This position offers the opportunity to work on impactful projects and engage directly with the ML community.
Must have:
  • Bachelor's degree or equivalent experience
  • 2+ years software development experience
  • 2+ years experience testing/launching software
  • 1+ year software design and architecture experience
  • Experience with data structures/algorithms
Good to have:
  • 4+ years experience in networking, storage, etc.
  • C, C++, Java, or Python programming experience
  • Experience building AI infrastructure
  • API design experience
  • Experience with low-level C++ and high-level Python
  • Knowledge of latest AI models and techniques
Perks:
  • Bonus
  • Equity
  • Benefits

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

Next Level Business Services - Java Developer

Next Level Business Services

San Jose, California, United States (On-Site)
6 Months ago
Next Level Business Services - Java Developer with Oracle SOA

Next Level Business Services

Cincinnati, Ohio, United States (On-Site)
7 Months ago
PwC - IN_Senior Associate _Java Developer _Data & Analytics _Advisory _PAN India

PwC

Kolkata, West Bengal, India (On-Site)
7 Months ago
NVIDIA - Data Engineer

NVIDIA

Taipei City, Taiwan (On-Site)
3 Months ago
Dream Sports - SDE 3 - Backend

Dream Sports

Mumbai, Maharashtra, India (On-Site)
9 Months ago
Easygo - Fullstack Software Development Engineer, Sportsbook

Easygo

Melbourne, Victoria, Australia (On-Site)
3 Months ago
Netflix - UI Engineer (L4/L5) - Enablement Apps

Netflix

Warsaw, Masovian Voivodeship, Poland (On-Site)
3 Months ago
Info Stretch - Senior Java Engineer

Info Stretch

Poland (On-Site)
6 Months ago
Sigma Software - Senior Java Developer

Sigma Software

Brasília, Federal District, Brazil (On-Site)
7 Months ago
Google - Staff Software Engineer, Mobile (Android), YouTube

Google

San Bruno, California, United States (On-Site)
3 Weeks ago

Get notifed when new similar jobs are uploaded

Similar Skill Jobs

The Walt Disney Company - Lead Software Engineer - Full-Stack

The Walt Disney Company

San Francisco, California, United States (On-Site)
1 Month ago
Crunchyroll - Staff Site Reliability Engineer

Crunchyroll

Mexico City, Mexico City, Mexico (On-Site)
6 Months ago
Actian - Sustenance Engineer - Actian Data Platform - Bangalore/Pune

Actian

Bengaluru, Karnataka, India (On-Site)
7 Months ago
ByteDance - Software Development Engineer Graduate (Distributed NoSQL Database Systems)

ByteDance

Seattle, Washington, United States (On-Site)
1 Month ago
Meta - Software Engineering Manager, Product Infrastructure

Meta

Menlo Park, California, United States (Remote)
6 Months ago
Microsoft - Member of Technical Staff, AI - Multimodal

Microsoft

(On-Site)
1 Month ago
Tencent - Data Engineer Intern

Tencent

Amsterdam, North Holland, Netherlands (On-Site)
3 Months ago
Microsoft - Software Engineer

Microsoft

Ho Chi Minh City, Ho Chi Minh City, Vietnam (On-Site)
1 Month ago
Google - Software Engineering Manager II, Infrastructure, Google Cloud

Google

Raleigh, North Carolina, United States (On-Site)
1 Month ago
NinjaVan - Senior Data Engineer

NinjaVan

Hyderabad, Telangana, India (On-Site)
7 Months ago

Get notifed when new similar jobs are uploaded

Jobs in Mountain View, California, United States

Google - Senior Software Engineer, Embedded Systems/Firmware, Pixel

Google

Mountain View, California, United States (On-Site)
1 Month ago
ByteDance - Software Engineer in Machine Learning Systems

ByteDance

Seattle, Washington, United States (On-Site)
6 Months ago
ByteDance - Software Engineer Graduate (Backend-Dev Infra)

ByteDance

San Jose, California, United States (On-Site)
1 Month ago
People Can Fly - Live Operations Technician

People Can Fly

Yonkers, New York, United States (Remote)
2 Months ago
The Walt Disney Company - News Editor

The Walt Disney Company

Glendale, California, United States (Hybrid)
1 Month ago
Daybreak Game Company LLC - QA Lead

Daybreak Game Company LLC

Renton, Washington, United States (Hybrid)
3 Months ago
ByteDance - TikTok Shop Fashion Category Manager SMB

ByteDance

New York, New York, United States (On-Site)
1 Month ago
The Walt Disney Company - Senior Manager, Safety

The Walt Disney Company

Anaheim, California, United States (On-Site)
3 Weeks ago
Patreon - Senior Backend Engineer, Creation

Patreon

New York, New York, United States (Hybrid)
1 Month ago
Netflix - Data Engineer (L4) - Security

Netflix

United States (Hybrid)
3 Months ago

Get notifed when new similar jobs are uploaded

Full Stack Development Jobs

Google - Senior Software Engineer, Full Stack, Chrome Tooling

Google

Warsaw, Masovian Voivodeship, Poland (On-Site)
3 Weeks ago
Google - Uber Technical Lead, Spaces Engineering

Google

Hyderabad, Telangana, India (On-Site)
1 Month ago
Google - Senior Staff Software Engineer, Accelerators Systems Software Architect

Google

Sunnyvale, California, United States (On-Site)
3 Weeks ago
Google - Software Engineer III, NetSight, Google Cloud

Google

Tel Aviv-Yafo, Tel Aviv District, Israel (On-Site)
3 Weeks ago
Google - Senior Software Engineer, Full Stack, Guided Support Experience

Google

Pittsburgh, Pennsylvania, United States (On-Site)
3 Weeks ago
Nagarro - Staff Engineer, Sharepoint

Nagarro

India (Remote)
7 Months ago
PwC - Senior Associate_Full Stack Developer_Data & Analytics_Advisory_PAN  India

PwC

Kolkata, West Bengal, India (On-Site)
7 Months ago
Google - Software Engineer, XProf

Google

Sunnyvale, California, United States (On-Site)
1 Month ago
ARHS - Fullstack Developer

ARHS

Liège, Wallonia, Belgium (On-Site)
7 Months ago
N-iX - Senior Node.js Engineer

N-iX

Ukraine (Remote)
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