NVIDIA Corporation

Senior Software Engineer, JAX

NVIDIA Corporation$184K — $356K *
US-AnywhereRemote in Santa Clara, CA
Enterprise Technology
5 - 7 years of experience
Job Overview by Ladders

Qualifications

  • BS in Computer Science, Computer Engineering or related field (or equivalent experience)
  • 6+ years of relevant experience
  • Proficient in C/C++ and Python programming
  • Experience with machine learning frameworks (e.g., PyTorch, TensorFlow)
  • Strong technical knowledge of CPU and GPU architectures and numeric libraries
  • Proven ability to develop customer-facing solutions
  • Excellent verbal and written communication skills

Responsibilities

  • Contribute significantly to NVIDIA's initiatives in JAX
  • Design and implement core components of JAX for peak performance on NVIDIA products
  • Collaborate with AI researchers to build robust future-proof models
  • Develop tools to enhance team efficiency in AI system development
  • Bridge numerical computing, simulation, and deep learning research with real-world applications

Benefits

  • Eligibility for equity
  • Diverse work environment
  • Commitment to equal opportunity employment
Full Job Description
NVIDIA is hiring senior engineers to develop its AI platform and more specifically its performance optimizations in deep learning frameworks using JAX, a tool that can differentiate between Python control code and NumPy code. Our goal is to deliver a polished, fast, modular and coordinated platform to handle data, training and analysis for a wide range of deep learning solutions using JAX. To achieve this and provide an excellent user experience, you have strong programming skills and excellent system design skills. You should have excellent communication and planning skills for this role. Come join us to build the next generation of tooling that will advance the numerical computing and machine learning research around the world!

What You Will Be Doing:
  • Play meaningful role in NVIDIA's effort in contributing to JAX.
  • Design and implement JAX core components and drive peak performance on NVIDIA products.
  • Work with AI applied researchers and leaders to build future-proof models
  • Build tools that will increase the efficiency of teams developing AI-based systems.
  • Work to bridge the gap between the latest in numerical computing, simulation and deep learning research and their applications in real world products.


What We Need To See:
  • BS in Computer Science or Computer Engineering or related field (or equivalent experience)
  • 6+ years relevant experience
  • C/C++ and Python programming
  • Experience with machine learning frameworks and their internals (e.g. PyTorch, TensorFlow, scikit-learn, etc.)
  • Proven ability developing customer-facing solutions, balancing feature requests and bugs.
  • Proven technical foundation in CPU and GPU architectures, numeric libraries, modular software design.
  • Highly motivated with excellent verbal and written communication skills.
  • Ability to work successfully with multi-functional teams, principles and architects. Coordinates effectively across organizational boundaries and geographies.


Ways To Stand Out From The Crowd:
  • Understanding of JAX, Autograd, tracing, code generation and DSL compilers and their design.
  • Understanding of deep learning training in distributed contexts: multi-GPU, multi-node, synchronous vs asynchronous.
  • Background with software shipping cycles (dev, deploy, release, CI).
  • Experience building distributed systems and services at large scale.


Your base salary will be determined based on your location, experience, and the pay of employees in similar positions. The base salary range is 184,000 USD - 287,500 USD for Level 4, and 224,000 USD - 356,500 USD for Level 5.

You will also be eligible for equity and benefits.

Applications for this job will be accepted at least until July 13, 2026.

This posting is for an existing vacancy.

NVIDIA uses AI tools in its recruiting processes.

About NVIDIA Corporation

Nvidia, a global leader in graphics, gaming, and AI technology, offers Nvidia careers and internship opportunities for those passionate about driving innovation in the tech industry. you'll find a company committed to growth, teamwork, and leadership in computer science and machine learning domains.

About Nvidia

A Pioneer in Technology and Innovation

Nvidia has cemented its reputation as a powerhouse in developing advanced graphics processing units (GPUs) and has significantly contributed to the gaming industry's evolution. Moreover, its foray into AI and machine learning has opened new frontiers in technology, making Nvidia a beacon of innovation and a desirable workplace for ambitious tech professionals.

Job Opportunities

Diverse Positions in a Dynamic Field

Nvidia is continuously on the lookout for talented individuals across various domains, including hardware and software engineering, product design, marketing, and sales. Employment opportunities at Nvidia are vast, catering to a wide range of expertise and career aspirations.

Employment in Hardware and Graphics

For those fascinated by the intricacies of hardware and graphics technology, Nvidia offers positions that sit at the forefront of gaming and computing advancements.

Growth in Machine Learning and AI

Nvidia's leadership in AI and machine learning has created numerous vacancies for specialists eager to contribute to groundbreaking projects.

Recruitment in Computer Science

With the constant demand for innovation, Nvidia's recruitment efforts focus on computer science experts capable of pushing the boundaries of what's possible.

Internship Program

Opening Doors to Future Innovators

Nvidia's internship program is designed to nurture the next generation of technology leaders, offering hands-on experience in a culture that celebrates creativity and teamwork.

Benefits and Culture

Interns at Nvidia enjoy a plethora of benefits, from competitive stipends to mentorship opportunities, all within an environment that values growth and learning.

Opportunities for Students

Whether you're an undergraduate, a master's student, or a Ph.D. candidate, Nvidia's internships provide a real-world glimpse into the tech industry, offering valuable experience in various technology fields.

Pathways to Full-Time Employment

Many interns have transitioned into full-time positions, marking the start of successful careers at Nvidia. The internship program is more than a stepping stone into the company; it’s an investment in the professional development of interns. The goal is to ensure that interns are well-equipped for future challenges.

Nvidia Careers: More Than Just a Job

Nvidia offers more than just a job to its employees; it provides a front-row seat on the journey into the future of technology. Nvidia stands as a pillar of innovation with its vast opportunities in hardware, graphics, gaming, machine learning, and computer science. Nvidia careers serve as a launching pad for talented workers who aim to redefine the technological landscape. Whether through full-time positions or internships, joining Nvidia means contributing to a legacy of breakthroughs and becoming part of a global community dedicated to pushing the boundaries of what's possible.
Learn more about NVIDIA Corporation
Size
22,473 employees
Market Cap
$350.4 billion
Industry
Net Income
$4.3 billion
Founded
1993
5 Year Trend
+31.3%
Revenue
$16.6 billion
NASDAQ

Similar Jobs

More Jobs at NVIDIA Corporation

More Enterprise Technology Jobs

Find similar Senior Software Engineer, JAX jobs: