Rochisha Agarwal

Mentors: Eric Giguère & Alex Pitchford

Background

I’m a physicist by training, with a solid grounding in quantum mechanics and computational methods, having completed my Bachelors and Masters of Physics from the Indian Institute of Technology Roorkee.

This summer Google Summer of Code gave me the opportunity to integrate the two powerful tools QuTiP and JAX , with QuTiP organization under the NUMFocus umbrella, with the goal of accelerating quantum simulations and enabling efficient gradient computations.

Objective

The primary objective of this project was to integrate JAX with QuTiP, leveraging JAX’s capabilities for automatic differentiation and just-in-time compilation to enhance the current QuTiP library. The end goal was to enable faster and more efficient optimization techniques in quantum simulations, particularly for solving complex quantum systems and computing gradients that are crucial for optimization routines.

Why QuTiP and JAX?

QuTiP Overview

QuTiP (Quantum Toolbox in Python) is a comprehensive software package widely used in the quantum computing community for simulating quantum systems, particularly open quantum systems. It offers a range of tools for solving various quantum mechanical problems, including the evolution of quantum states, computing observables, and performing advanced quantum control simulations. QuTiP is highly valued for its flexibility, ease of use, and the extensive library of functions it provides for quantum mechanics research.

JAX Overview

JAX is a modern computational framework that offers powerful tools for automatic differentiation and just-in-time (JIT) compilation. Originally developed by Google for machine learning research, JAX has become popular in various scientific computing fields due to its ability to efficiently compute gradients, optimize functions, and leverage hardware acceleration (e.g., GPUs and TPUs). It seamlessly integrates with NumPy, making it an attractive choice for enhancing performance in numerical simulations.

The Intersection

Integrating JAX with QuTiP brings the best of both worlds: QuTiP’s rich quantum mechanics toolkit and JAX’s computational power. The primary motivation behind this integration is to accelerate quantum control simulations, particularly those involving optimization routines where gradient computation is crucial. By leveraging JAX’s grad function and JIT compilation, we can achieve automatic differentiation, useful for variety of simulations.

What is Automatic Differentiation (jax.grad)?

Automatic differentiation is a technique used to compute derivatives of functions efficiently and accurately. Unlike numerical differentiation, which approximates derivatives by perturbing inputs, or symbolic differentiation, which involves manipulating algebraic expressions, automatic differentiation breaks down the computation of a function into a series of elementary operations and applies the chain rule to each operation.

Why is it Useful?

• Machine Learning: In machine learning, particularly in training neural networks, gradients are used to optimize the model’s parameters by minimizing a loss function. The ability to compute gradients efficiently is crucial for gradient descent and its variants, which are at the core of most machine learning algorithms. JAX’s grad function makes it easy to implement these optimization routines with high performance.