Error in optimization with gate_split

[rafaeleb]

Hello all,

I am currently implementing an optimization procedure for an MPS, where the loss function is written as

`def loss(psi):

all_terms = [] 
#kinetic terms
for i in range(l):
    phi = cp.copy(psi)
    all_terms.append(phi.gate(jnp.array(0.5 * D2), l, contract = True))

#potential terms
for i in range(l):
    phi = cp.copy(psi)
    all_terms.append(phi.gate(jnp.array(0.5 * X2), l, contract = True))

#two-body interactions
for i in range(l-1):
    phi = cp.copy(psi)
    all_terms.append(phi.gate_split(jnp.array(g*Xx), (i, i+1)))

final_psi = np.sum(np.array(all_terms))
return psi.H @ final_psi #returns the expectation value of the energy <psi|H|psi>`

works works for an individual calculation of the loss, but when running the optimizer

tnopt = TNOptimizer(psi, loss_fn = loss, norm_fn = norm_fn, optimizer = "adam", callback = callback, autodiff_backend = 'jax')
I get an error from JAX (just pasting the last lines):

ConcretizationTypeError: Abstract tracer value encountered where concrete value is expected: Traced<ShapedArray(bool[])>with<DynamicJaxprTrace(level=1/0)> The problem arose with the boolfunction.

I could not find the problematic function but it's related to the application of .gate_split, since the rest of the loss function does not raise errors. Any way I can solve this issue?

Thanks,

Rafael

[jcmgray]

Hi @rafaeleb, you might be able to fix this by using a fixed max_bond and cutoff=0.0 in the gate_split function. The problem currently is that the shape of the traced computational graph depends on the cutoff. Also, you should be able to use torch which doesn't require a fixed computational graph.

[rafaeleb]

Hi Johnnie,

Thanks for the reply. The addition of the compression options did not fix the problem, but it looks like it can be avoided in my particular case by not using gate_split at all (the gates can be decomposed from the start). But it would be nice to know if the gate_split function can be modified is some way to agree with the jax optimizer backend.

BR