DMRG with jax or torch

[yuxuanzhang1995]

What is your issue?

Hi quimb support team, I would like to train a neural network with a DMRG solver as a sub-routine. Does quimb currently support DMRG with torch or jax? Will the AD for SVD work there? Thanks!

[jcmgray]

Hi @yuxuanzhang1995, currently it likely will not work. This is because whilst most of the array operations quimb performs are backend agnostic, the DMRG routine specifically uses a TNLinearOperator to represent the local hamiltonian operator, and this eigensystem is solved iteratively using scipy or slepc, both of which require numpy arrays. To fix this one would need:

1. to modify TNLinearOperator so that the input and output vector can be arbitrary backend vectors
2. add a backend agnostic iterative eigensolver, e.g. like conjugate_gradient, or dispatch to the jax or torch ones.