Senior Software Engineer, JAX

• 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
• Bridge the gap between numerical computing, simulation and deep learning research and real-world product applications

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