Callbacks and GPUs

While Catalyst aims to support all classical processing functionality as provided by JAX, there are sometimes cases where you may need to perform a host callback to execute arbitrary Python-compatible code. This may include use-cases such as:

• runtime debugging and logging,

• executing classical subroutines on accelerators such as GPUs or TPUs, or

• incorporating non-JAX compatible classical subroutines within a larger QJIT workflow.

Catalyst supports all of these via a collection of callback functions.

Overview

Catalyst provides several callback functions:

• callback() supports callbacks of functions with no return values. This makes it an easy entry point for debugging, for example via printing or logging at runtime.

• pure_callback() supports callbacks of pure functions. That is, functions with no side-effects that accept parameters and return values. However, the return type and shape of the function must be known in advance, and is provided as a type signature.

  Note that to use pure_callback() within functions that are being differentiated, a custom VJP rule must be defined so that the Catalyst compiler knows how to differentiate the callback. This can be done via the pure_callback.fwd and pure_callback.bwd methods. See the pure_callback() documentation for more details.

• accelerate() is similar to pure_callback() above, but is designed to work only with functions that are jax.jit compatible. As a result of this restriction, return types do not have to be provided upfront, and support is provided for executing these callbacks directly on classical accelerators such as GPUs and TPUs.

In addition, print(), a convenient wrapper around the Python print function, is provided for runtime printing support.

Callbacks to arbitrary Python

When coming across functionality that is not yet supported by Catalyst, such as functions like scipy.integrate.simpson, Python callbacks can be used to call arbitrary Python code within a qjit-compiled function, as long as the return shape and type is known:

import scipy as sp

@pure_callback
def simpson(x, y) -> float:
    return sp.integrate.simpson(y, x=x)

@qjit
def integrate_xsq(a, b):
    x = jnp.linspace(a, b, 100)
    return simpson(x, x ** 2)

>>> integrate_xsq(-1, 1)
Array(0.66666667, dtype=float64)
>>> integrate_xsq(-1, 2)
Array(3., dtype=float64)

Please see the docstring of pure_callback() for more details, including how to define vector-Jacobian product (VJP) rules for autodifferentiation, and for specifying the return-type of vector-valued functions.

Callbacks to JIT-compatible code

If a function is JIT-compatible, then accelerate() can be used, negating the need to manually provide return shape and dtype information:

@qjit
def fn(x):
    x = jnp.sin(x)
    y = catalyst.accelerate(jnp.fft.fft)(x)
    return jnp.sum(y)

>>> x = np.array([1.0, 2.0, 1.0, -1.0, 1.5])
>>> fn(x)
Array(4.20735492+0.j, dtype=complex128)

Accelerated functions also fully support autodifferentiation with grad(), jacobian(), and other Catalyst differentiation functions, without needing to specify VJP rules manually:

@qjit
@grad
def f(x):
    expm = catalyst.accelerate(jax.scipy.linalg.expm)
    return jnp.sum(expm(jnp.sin(x)) ** 2)

>>> x = jnp.array([[0.1, 0.2], [0.3, 0.4]])
>>> f(x)
Array([[2.80120452, 1.67518663],
       [1.61605839, 4.42856163]], dtype=float64)

Accelerator (GPU and TPU) support

accelerate() can also be used to execute classical subroutines on classical accelerators such as GPUs and TPUs:

@accelerate(dev=jax.devices("gpu")[0])
def classical_fn(x):
    return jnp.sin(x) ** 2

@qjit
def hybrid_fn(x):
    y = classical_fn(jnp.sqrt(x)) # will be executed on a GPU
    return jnp.cos(y)

dev/callbacks

 

Download Python script

 

Download Notebook

 

View on GitHub