Talent.com
Senior Software Engineer, JAX
Senior Software Engineer, JAXNVIDIA • Remote, CA, US
Senior Software Engineer, JAX

Senior Software Engineer, JAX

NVIDIA • Remote, CA, US
30+ days ago
Job type
  • Full-time
  • Remote
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)

  • 5+ 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.

The base salary range is 180,000 USD - 339,250 USD. Your base salary will be determined based on your location, experience, and the pay of employees in similar positions.

You will also be eligible for equity and . NVIDIA accepts applications on an ongoing basis.

NVIDIA is committed to fostering a diverse work environment and proud to be an equal opportunity employer. As we highly value diversity in our current and future employees, we do not discriminate (including in our hiring and promotion practices) on the basis of race, religion, color, national origin, gender, gender expression, sexual orientation, age, marital status, veteran status, disability status or any other characteristic protected by law.

Create a job alert for this search

Senior Software Engineer, JAX • Remote, CA, US

Similar jobs

Remote-Senior Software Engineer-Intuit

S M Software Solutions IncPiedra, CA, United States
Remote
Full-time

Job description :We are looking for a Senior Software Engineer specializing in Retrieval-Augmented Generation (RAG) systems, with experience in large language models (LLMs), vector databases, and c...Show more

 • Promoted

Jira Engineer

Unicorn Technologies LLCCA, United States
Full-time
Quick Apply

MessageBody">Job Title: Jira Engineer Location: San Jose, CA (hybrid 3 days a week) The customer's JIRA instance is highly complex and heavily customized.This role goes beyond standard JIRA adminis...Show more

Automation Engineer

Sumeru SolutionsCA, United States
Full-time
Quick Apply

Job description We are seeking an experienced Automation Engineer to design implement and maintain robust automation frameworks for web mobile and OTT appl...Show more

Earn $60K–$115K While Job Searching — Become a Surrogate on Your Timeline

Alpha FidelisEast Tulare Villa, CA, US
Full-time +1

Earn While You Job Search — Become a Surrogate ✨.Searching for your next opportunity while supporting your family isn’t easy.What if you could earn meaningful income.Why This Opportunity Stands Out...Show more

 • Promoted

Sales Associate

BootBarn, Inc.Visalia, CA, United States
Full-time

Our Sales Partners engage with customers and create an exceptional in-store experience.They embody Boot Barn's Mission, Vision and Values in their words and actions while demonstrating selling and ...Show more

 • Promoted

Senior Full Stack Engineer

FirestormRemote, California, United States,
Remote
Permanent

Firestorm is building the next generation of uncrewed aircraft and the advanced manufacturing systems that deliver them at speed.The Software Integration & Operations department owns the software l...Show more

Principal Software Engineer

pubGENIUSCA, US
Remote
Full-time
Quick Apply

We build websites and apps for clients across the US and Europe, specializing in decentralized finance (DeFi, crypto, NFT, blockchain), fintech, and healthcare, with occasional work in e-commerce a...Show more

Java Software Engineer

OpenkyberCA, United States
Full-time
Quick Apply

Location: Greenwood Village, CO Salary: $70.USD Hourly Job Title: Software Development Engineer - Java And.Net Location: Greenwood Village, CO / Berkeley Heights, NJ / Omaha, NE (Onsite) Duration: ...Show more

Java QA Automation Engineer

OpenkyberCA, United States
Temporary
Quick Apply

Greetings, Very Urgent fill , NO H1s Please.Job Title: QA Automation Engineer / SDET Location : Phoenix, AZ ( Local Needed ) Duration: 12+ Months Contract Job Summary: We are seeking a highly skill...Show more

Area Supervisor

Ross Stores, IncTulare, CA, United States
Full-time

Our values start with our people, join a team that values you! Bring your talents to Ross, our leading off-price retail chain with over 2,200 stores, and a strong track record of success and growth...Show more

 • Promoted

Senior Principal Software Engineer with Test Equipment

RaytheonCA, United States
Temporary

US-AZ-TUCSON-801 ~ 1151 E Hermans Rd ~ BLDG 801 (External Site).Person, or Immigration Status Requirements:.At Raytheon, the foundation of everything we do is rooted in our values and a higher call...Show more

 • Promoted

Sr. Full stack Engineer - Remote

Two95 International Inc.CA, US
Remote
Full-time
Quick Apply

Create clean, maintainable, and scalable, and well-tested code.Create elegant web-based user interfaces and reporting dashboards.Work with other team members to devise the best possible technical s...Show more

Earn $60K-$115K While Job Searching - Become a Surrogate on Your Timeline

Alpha FidelisVisalia, CA, United States
Permanent

Earn While You Job Search — Become a Surrogate.Interested in this role You can find all the relevant information in the description below.Searching for your next opportunity while supporting your f...Show more

 • Promoted

AWS API Engineer

OpenkyberCA, United States
Temporary
Quick Apply

Job Title: Java Microservice Architect Location: Santa Clara, CA (Onsite) Job Type: Contract (Long-Term) Job Description: We are seeking an experienced Java Microservice Archit...Show more

Senior Software Engineer

RaytheonCA, United States
Full-time

US-CA-EL SEGUNDO-E04 ~ 2006 E El Segundo Blvd ~ BLDG E04.Person, or Immigration Status Requirements:.Active and existing security clearance required on day 1.At Raytheon, the foundation of everythi...Show more

 • Promoted • New!

Full time Software Engineers in San Diego, CA

Maania Consultancy ServicesCalifornia, CA, US
Full-time
Quick Apply

Our client is looking for a Software Engineers (Python, Node.If you're interested please share your updated resume along with your yearly expected salary Role:.Software Engineers (Python, Node.C) L...Show more

Shift Supervisor - Visalia (58188)

Inter-Con Security Systems, Inc.Visalia, CA, United States
Full-time +1

Overview Founded in 1973, Inter-Con Security Systems, Inc.US-owned security company, providing integrated security solutions to government and commercial customers on four continents.Inter-Con rema...Show more

 • Promoted

Senior Software Engineer with Test Equipment - 4th Shift

RaytheonCA, United States
Temporary

US-AZ-TUCSON-801 ~ 1151 E Hermans Rd ~ BLDG 801 (External Site).Person, or Immigration Status Requirements:.At Raytheon, the foundation of everything we do is rooted in our values and a higher call...Show more