Has anybody tried using OQuPy with a GPU?

[gefux]

We currently use TensorNetwork as a backend in OQuPy, which in turn can use numpy/pytorch/tensorflow as a backend. Has anybody tried using OQuPy with tensorflow on a GPU?

[Sampreet]

I tried to run OQuPy on the GPU, although not with PyTorch/TensorFlow, but with JAX. Since TensorNetwork's Node requires its tensor parameter to be passed with the same backend, and given that most of OQuPy's modules use NumPy explicitly to create these nodes, I tried to implement a GPU-based numerical backend by switching from numpy.ndarray to jax.numpy.ndarray dynamically in my fork (very similar to what has been recently implemented in QuTiP to support auto-differentiation for the solvers in its JAX backend).

Updating the Modules to Support JAX

The oqupy.config module is modified as:

import numpy as default_numpy
import scipy as default_scipy
NUMERICAL_BACKEND_NUMPY = default_numpy
NUMERICAL_BACKEND_LINALG = default_scipy.linalg
NUMPY_DTYPE_COMPLEX = default_numpy.complex128   # earlier NpDtype
NUMPY_DTYPE_REAL = default_numpy.float64         # earlier NpDtypeReal

The specific choice of scipy.linalg instead of scipy is because jax.scipy.integrate doesn't yet have implementations of dblquad, quad and quad_vec as required by some of OQuPy's modules.
A new oqupy.backends.numerical_backend module now takes care of the switching. The coarse contents of the module are:

class NumPy:
    @property
    def backend(self) -> default_numpy:
        return oqupy.config.NUMERICAL_BACKEND_NUMPY
    def __getattr__(self, name):
        backend = object.__getattribute__(self, 'backend')
        return getattr(backend, name)
    # additional overridden methods
class SciPyLinAlg:
    # same as above for ``default_scipy.linalg``
np = Numpy()
sl = SciPyLinAlg()

The calls for numpy in all of OQuPy's modules would be changed from:

import numpy as np
import scipy.linalg as linalg
from oqupy.config import NpDtype, NpDtypeReal

to:

from oqupy.backends.numerical_backend import np, sl
# change corresponding occurrences of NpDtype and NpDtypeReal

With a few more changes to override in-place updates and random number generations in oqupy.backends.numerical_backend, and fix non-array inputs in other OQuPy modules, I tried the example implementations and tests of OQuPy.

Running an Example

Any script using the JAX backend would require the following lines:

# packages
import jax
import jax.numpy as jnp
import jax.scipy.linalg as jsl
import oqupy.config as oc
import tensornetwork as tn
from oqupy.backends.numerical_backend import np, sl
# configuration
jax.config.update('jax_enable_x64', True)
oc.NUMERICAL_BACKEND_NUMPY = jnp
oc.NUMPY_DTYPE_COMPLEX = jnp.complex128
oc.NUMPY_DTYPE_REAL = jnp.float64
oc.NUMERICAL_BACKEND_LINALG = jsl
tn.set_default_backend('jax')

The runtimes (without jitted functions) corresponding to $\alpha = 0.1$ in Fig. 4(b) of arXiv:2406.16650v1 obtained using my Intel i7 8700K and NVIDIA GeForce RTX 3060Ti are:

Backend	Runtime (s)
NumPy	~12
JAX (CPU)	~88
JAX (GPU)	~430

Further optimizations (for-looped updates, lists -> arrays) only reduced the runtimes by a couple of seconds.

Issues

The tests involving the NumPy backend went smoothly! However, the JAX-backend tests had multiple issues. Before attempting to resolve them, I wanted to know if any such change would break the flow of the package. A couple of these issues are mentioned below:

Immutability of jax.numpy Arrays

The add_singleton function in oqupy.utils is often called by other modules (primarily via oqupy.backends.pt_tempo_backend) to update MPOs and MPSs. This function updates the shape of the array by adding an additional dimension to it at a specified index. Whereas NumPy supports altering the shape attribute, JAX's ArrayImpl datatypes are immutable and as such shape doesn't have a setter method leading to a recurrent error:

AttributeError: property 'shape' of 'ArrayImpl' object has no setter

Similar mutations can be found in oqupy.gradient, oqupy.mps_mpo and oqupy.system_dynamics modules.

In this case, would reshaping the array instead of updating its shape break the functionality?

Vectorization Signatures

Modules such as oqupy.bath_correlations and oqupy.system use the vectorize method of NumPy without explicitly mentioning the signatures of the input and output. This raises issues with JAX when the parameters contain non-scalar inputs, such as the Hamiltonian or the Lindblad operators. Depending on the type of implementation, the signatures may be added for each vectorized function. For example, the Hamiltonians for TimeDependentSystem and TimeDependentSystemWithField will have the signatures ()->(m,m) and (),()->(m,m) respectively (kindly correct me if I have made a mistake here).

Row Degeneracy Checks

Unlike the numpy.unique method, the jax.numpy.unique method returns array indices with the same shape as the input array even with the axis parameter. Although this can be resolved by using the flatten method, I am not sure if it will introduce additional errors.

It would be great if someone could provide suggestions for the mentioned issues or suggest workarounds (if any) for the same. Also, kindly let me know if I should raise an Issue mentioning these.

[piperfw]

Hi @Sampreet, as per our email discussion thanks for the excellent work on this and for so clearly laying out relevant points of discussion for us. Your contribution certainly has enough structure for us to move to the Issue section! Continued in #142.