Using GPUs for the compression

[benjaminlanthier]

Hi Johnnie,

I have been using the compress_all() function in my tensor networks for some time now using numpy arrays in the quimb.tensor.Tensor() objects' data and I would like to use the GPUs for it in order to push my computations further. I know that quimb uses autoray, so the data backend in the Tensor objects should be the one used for the compression. Is that right? If it is, it would mean that if I was to use jax.numpy for my tensors' data and that a GPU is available, the decomposition function would operate on the GPU and if none was available, it would run on the CPU ?

[jcmgray]

Hi @benjaminlanthier, yes that is the essential idea - the array operators are dispatched based on the data they are supplied with, and indeed jax defaults to gpu if available else falls back to cpu.

Note that to get good performance out of jax has historically required jit-compiling the functions - it less works well as a dynamic array library in the vein of numpy or torch.