Senior Software Engineer, JAX

Requirements:

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

What you'll 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.

Nice to Haves:

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

Perks and Benefits:

• Your base salary will be determined based on your location, experience, and the pay of employees in similar positions.
• For Poland: The base salary range is 292,500 PLN - 507,000 PLN for Level 4, and 375,000 PLN - 650,000 PLN for Level 5.