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 role involves delivering a polished, fast, modular, and coordinated platform for data, training, and analysis in deep learning solutions.

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

Nice to have:

• Understanding of JAX, Autograd, tracing, code generation, 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.

• 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 numerical computing, simulation, deep learning research, and real world products.