Senior Software Engineer, JAX

NVIDIA is hiring senior engineers to develop its AI platform, focusing on performance optimizations in deep learning frameworks using JAX. The goal is to deliver a polished, fast, modular and coordinated platform for data, training and analysis across a wide range of deep learning solutions using JAX. Join us to build the next generation of tooling that will advance numerical computing and machine learning research worldwide.

• Play a 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 to increase the efficiency of teams developing AI-based systems.
• Bridge the gap between the latest in numerical computing, simulation and deep learning research and their applications in real world products.

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

NVIDIA develops AI platforms and deep learning solutions, advancing numerical computing and machine learning research globally.

Department: Engineer, Tech SW. Collaboration with AI applied researchers and leaders. Work successfully with multi-functional teams, principles and architects. Coordinates across organizational boundaries and geographies.