Learn how to write blazingly fast and modern numerical code in Python by leveraging just-in-time (JIT) compilation for CPU’s and GPU’s and how to scale computations across multiple machines. Tailored for data scientists, researchers, and Python enthusiasts, this session will focus on practical applications of JAX, Numba, and Ray to optimise and parallelise your numerical code. These techniques have applications across multiple fields like machine learning, numerical simulations, or engineering applications. You can already forget low level languages like C/C++, Fortran or similar.

What You’ll Learn

1. Optimising and Accelerating Numerical Code with JAX and Numba:
   • Converting NumPy code to JAX/Numba.
   • JIT compilation for speeding up Python functions.
   • Automatic differentiation (JAX).
   • GPU/CPU acceleration and vectorisation.
   • Main limitations of JAX and Numba.
   • Hands-on exercises combining JAX and Numba to optimise loops, linear algebra, and scientific calculations.
2. Scaling Workflows with Ray:
   • Introduction to Ray: Parallel, out-of core and distributed computing made simple.
   • Managing distributed tasks, scaling workloads across CPUs/GPUs, and handling large datasets.
   • Integrating JAX and Numba in Ray-powered workflows for end-to-end acceleration.

Who Should Attend?

Participants should have intermediate Python programming skills and basic familiarity with NumPy and linear algebra. No prior experience with JAX, Numba, or Ray is necessary.

Structure

• Hour 1: Foundations of JAX and Numba – JIT compilation, GPU acceleration, and their respective strengths.
• Hour 2: Hands-on optimisation: Writing high-performance numerical code with JAX and Numba.
• Hour 3: Scaling computations with Ray: Parallelising JAX/Numba workflows across machines.

Takeaways

• Understand when and how to use JAX and Numba for accelerating numerical Python code.
• Write scalable, high-performance workflows using Ray.
• Leave with ready-to-use examples and insights to apply in their own projects.