TF optimizer for complex-valued PEPS

[gabrielwoolls]

Hi @jcmgray, I've been having some trouble using the TNOptimizer for complex-valued networks using tensorflow backend. I'm attaching an example where optimizing a small PEPS crashes if we try to do boundary MPS contraction.

import quimb as qu
import quimb.tensor as qtn
from quimb.tensor.optimize import TNOptimizer

LX, LY = 2, 2
DTYPE = 'complex128'

def state_energy(psi, hterms, vterms, **opts):

    he = psi.compute_local_expectation(
        hterms, normalized=False, contract_optimize='random-greedy',**opts)

    ve = psi.compute_local_expectation(
        vterms, normalized=False, contract_optimize='random-greedy',**opts)        

    return he + ve

def normalize_state(psi):
    return psi.normalize()

def main():
    Hij = qu.ham_heis(2)

    peps = qtn.PEPS.rand(LX, LY, bond_dim=2, dtype=DTYPE)

    hterms = {coos: Hij for coos in peps.gen_horizontal_bond_coos()}
    vterms = {coos: Hij for coos in peps.gen_vertical_bond_coos()}

    compute_expec_opts = dict(
                    cutoff=2e-3, 
                    max_bond=9, 
                    contract_optimize='random-greedy')

    optmzr = TNOptimizer(
        peps, 
        loss_fn=state_energy,
        norm_fn=normalize_state,
        loss_constants={'hterms': hterms,
                        'vterms': vterms},
        loss_kwargs=   {'opts': compute_expec_opts},
        autodiff_backend='tensorflow',
    )

    peps_opt = optmzr.optimize(1)
    return peps_opt

if __name__ == '__main__':
    main()

It seems like the tensor_split method somehow uses a double dtype even if the original PEPS is set to be complex, which causes TF to get confused.

File "C:\Users\gabri\miniconda3\envs\encoding_tns\lib\site-packages\quimb\tensor\tensor_2d.py", line 2445, in normalize
    nfact = norm.contract_boundary(**boundary_contract_opts)
  File "C:\Users\gabri\miniconda3\envs\encoding_tns\lib\site-packages\quimb\tensor\tensor_2d.py", line 1235, in contract_boundary
    **boundary_contract_opts,
  File "C:\Users\gabri\miniconda3\envs\encoding_tns\lib\site-packages\quimb\tensor\tensor_2d.py", line 671, in contract_boundary_from_bottom
    compress_sweep=compress_sweep, **compress_opts)
  File "C:\Users\gabri\miniconda3\envs\encoding_tns\lib\site-packages\quimb\tensor\tensor_2d.py", line 608, in _contract_boundary_from_bottom_multi
    **compress_opts)
  File "C:\Users\gabri\miniconda3\envs\encoding_tns\lib\site-packages\quimb\tensor\tensor_2d.py", line 575, in _contract_boundary_from_bottom_single
    yrange=yrange, **compress_opts)
  File "C:\Users\gabri\miniconda3\envs\encoding_tns\lib\site-packages\quimb\tensor\tensor_2d.py", line 457, in compress_row
    self.compress_between((i, j), (i, j - 1), **compress_opts)
  File "C:\Users\gabri\miniconda3\envs\encoding_tns\lib\site-packages\quimb\tensor\tensor_core.py", line 3328, in compress_between
    tid1, tid2, canonize_distance=canonize_distance, **compress_opts)
  File "C:\Users\gabri\miniconda3\envs\encoding_tns\lib\site-packages\quimb\tensor\tensor_core.py", line 3288, in _compress_between_tids
    tensor_compress_bond(Tl, Tr, **compress_opts)
  File "C:\Users\gabri\miniconda3\envs\encoding_tns\lib\site-packages\quimb\tensor\tensor_core.py", line 757, in tensor_compress_bond
    absorb=absorb, **compress_opts)
  File "C:\Users\gabri\miniconda3\envs\encoding_tns\lib\site-packages\quimb\tensor\tensor_core.py", line 1482, in split
    return tensor_split(self, *args, **kwargs)
  File "C:\Users\gabri\miniconda3\envs\encoding_tns\lib\site-packages\quimb\tensor\tensor_core.py", line 647, in tensor_split
    left, s, right = split_fn(array, **opts)
  File "C:\Users\gabri\miniconda3\envs\encoding_tns\lib\site-packages\quimb\tensor\decomp.py", line 217, in _svd
    max_bond, absorb, renorm)
  File "C:\Users\gabri\miniconda3\envs\encoding_tns\lib\site-packages\quimb\tensor\decomp.py", line 165, in _trim_and_renorm_SVD
    VH = VH * reshape(s, (-1, 1))
  File "C:\Users\gabri\miniconda3\envs\encoding_tns\lib\site-packages\tensorflow_core\python\ops\math_ops.py", line 902, in binary_op_wrapper
    return gen_math_ops.mul(x, y, name=name)
  File "C:\Users\gabri\miniconda3\envs\encoding_tns\lib\site-packages\tensorflow_core\python\ops\gen_math_ops.py", line 6121, in mul
    _ops.raise_from_not_ok_status(e, name)
  File "C:\Users\gabri\miniconda3\envs\encoding_tns\lib\site-packages\tensorflow_core\python\framework\ops.py", line 6606, in raise_from_not_ok_status
    six.raise_from(core._status_to_exception(e.code, message), None)
  File "<string>", line 3, in raise_from
tensorflow.python.framework.errors_impl.InvalidArgumentError: cannot compute Mul as input #1(zero-based) was expected to be a complex128 tensor but is a double tensor [Op:Mul] name: mul/

Is there something simple I'm missing?

This is using the tensor_2d branch for the PEPS stuff of course.

[jcmgray]

Yeah tensorflow is kinda picky about dtype, which is annoying, like here, because:

• tensorflow.linalg.svd returns (complex, real, complex)
• then doesn't let you multiply the singular values in...

I'll add an explicit check for tensorflow and complex, unless I think of something else cleaner.

(By the way you might also have to use cutoff=0.0 for autodiff, as none of the compilers/autodiff computational graphs like dynamic shapes.)

[jcmgray]

Some other issues:

• easy to fix - you need to use qu.ham_heis(2).astype(DTYPE)

Slightly harder:

• tensorflow doesn't support autodiff through the QR for complex dtypes. One would need to switch all the canonize methods everywhere to use 'svd' rather than 'qr'. In principle that's just changing some method='svd' kwargs, but I'm not totally sure if it can be controlled all the way from compute_local_expectation at the moment.

[jcmgray]

I've fixed svd for tensorflow and complex, but the 'tf + QR + complex + autodiff' problem remains.

It may be of interest that the latest (dev) versions of torch now support complex data - torch generally is more convenient to work with auto-diff wise as it doesn't require (slow and badly scaling) compilation of the computational graph to work pretty efficiently (unlike jax and tensorflow).

[jcmgray]

Ok so one way to skirt around the 'tf + QR + complex + autodiff' problem is override the autoray lookup of linalg.qr for tensorflow:

import autoray

def tf_qr(x):
    U, s, VH = autoray.do('linalg.svd', x)

    dtype = autoray.get_dtype_name(U)
    if 'complex' in dtype:
        s = autoray.astype(s, dtype)

    Q = U
    R = autoray.reshape(s, (-1, 1)) * VH

    return Q, R

autoray.register_function('tensorflow', 'linalg.qr', tf_qr)

The following now works for me:

import quimb as qu
import quimb.tensor as qtn
from quimb.tensor.optimize import TNOptimizer

LX, LY = 4, 4
DTYPE = 'complex128'

autodiff_backend = 'tensorflow'
autodiff_backend_opts = {'experimental_compile': True}

def state_energy(psi, hterms, vterms, **opts):

    he = psi.compute_local_expectation(
        hterms, normalized=True, **opts)

    ve = psi.compute_local_expectation(
        vterms, normalized=True, **opts)        

    return autoray.do('real', (he + ve))

def normalize_state(psi):
    return psi.normalize()

Hij = qu.ham_heis(2).astype(DTYPE)

peps = qtn.PEPS.rand(LX, LY, bond_dim=2, dtype=DTYPE)

hterms = {coos: Hij for coos in peps.gen_horizontal_bond_coos()}
vterms = {coos: Hij for coos in peps.gen_vertical_bond_coos()}

compute_expec_opts = dict(
                cutoff=2e-3, 
                max_bond=9, 
                contract_optimize='random-greedy')

Hij = qu.ham_heis(2).astype(DTYPE)

peps = qtn.PEPS.rand(LX, LY, bond_dim=2, dtype=DTYPE)

hterms = {coos: Hij for coos in peps.gen_horizontal_bond_coos()}
vterms = {coos: Hij for coos in peps.gen_vertical_bond_coos()}

compute_expec_opts = dict(
                cutoff=0.0, 
                max_bond=9, 
                contract_optimize='random-greedy')

optmzr = TNOptimizer(
    peps, 
    loss_fn=state_energy,
#         norm_fn=normalize_state,
    loss_constants={'hterms': hterms,
                    'vterms': vterms},
    loss_kwargs=compute_expec_opts,  # these weren't being supplied quite right
    autodiff_backend=autodiff_backend,
    **autodiff_backend_opts,
)

peps_opt = optmzr.optimize(1000)