Applying gates with jnp index

[haimengzhao]

Issue Description

Hi, I'm trying to apply a gate with the qubit index specified by an integer i from jax.numpy array. But I get ValueError: Illegal index specification. I think it's because isinstance(i, int) returns False.

I want to do this because I need to jit, and forcing type conversion using int(i) will make it unjittable.

How to Reproduce

import jax.numpy as jnp
import tensorcircuit as tc

c = tc.Circuit(1)
ind = jnp.array([0], dtype=int)
c.X(ind[0])

Error Output

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[139], line 3
      1 c = tc.Circuit(1)
      2 ind = jnp.array([0], dtype=int)
----> 3 c.X(ind[0])

File [~/miniconda3/lib/python3.10/site-packages/tensorcircuit/abstractcircuit.py:196](~/miniconda3/lib/python3.10/site-packages/tensorcircuit/abstractcircuit.py:196), in AbstractCircuit.apply_general_gate_delayed..apply_list(self, *index, **kws)
    194         apply(self, *ind, **kws)
    195 else:
--> 196     raise ValueError("Illegal index specification")

ValueError: Illegal index specification

Environment Context

OS info: macOS-10.16-x86_64-i386-64bit
Python version: 3.10.10
Numpy version: 1.24.4
Scipy version: 1.11.0
Pandas version: 2.0.3
TensorNetwork version: 0.4.6
Cotengra is not installed
TensorFlow is not installed
Jax version: 0.4.13
Jax installation doesn't support GPU
JaxLib version: 0.4.13
PyTorch is not installed
Cupy is not installed
Qiskit version: 0.25.0
Cirq version: 1.2.0
TensorCircuit version 0.10.0

[refraction-ray]

Good question, since the topology of the circuit is a static information (shape information of the whole tensor network for the circuit), it is unable to jit in general even if you can use jax integer as the index number. However, in many interesting cases, depends on your specific use cases, you can find some way to smartly encode the problem and enable jit. So if you like, you can share your exact use case so that I can figure out whether the problem under simulation is jittable or not.

[haimengzhao]

Hi @refraction-ray , thanks for the timely reply! I see why we must treat the topology as static. I'm okay with that. But since I'm inputing the topology as a list of indices, and list is non-hashable and therefore cannot be used as a static argument in jax.jit. Do you have any suggestions on how to work around on that? For example, my code is like (arch for architecture, a list of (i, j) specifying the positions to apply two qubit gates):

def apply_circuit(c, arch, unitary):
    assert len(arch) == len(unitary)
    for i in range(len(arch)):
        c.any(int(arch[i][0]), int(arch[i][1]), unitary=unitary[i])
    return c

Many thanks in advance!

[refraction-ray]

# Apply X on given indices while jittable

n = 6
K = tc.set_backend("jax")

cand = K.stack([K.eye(2), tc.gates.pauli_gates[1].tensor])

def apply_gates(c, inds):
    inds = K.cast(inds, "int32")
    mpos = []
    for i in range(n):
        mpos.append(tc.gates.Gate(cand[inds[i]]))
    gate_applied = tc.quantum.QuOperator([m[0] for m in mpos], [m[1] for m in mpos])
    c.mpo(*range(n), mpo=gate_applied)
    return c

@K.jit
def cal_loss(inds):
    c = tc.Circuit(n)
    c = apply_gates(c, inds)
    for i in range(n-1):
        c.cx(i, i+1)
    return K.real(c.expectation_ps(z=[1]))

cal_loss(K.convert_to_tensor([0, 1, 0, 0, 1, 0]))

A smart way to around I just think of that might help, feel free to ask more if this paradigm not fits

[haimengzhao]

Thanks @refraction-ray , this works well! I'm wondering how to generalize this to 2-qubit gates. I think we can use [0, 0, 0, 1, 2] to specify indices. But it's not clear how to consturct the mpos list while kepping cand an array. I tried the following

def apply_circuit(c, arch, unitary):
    assert len(arch) == len(unitary)
    iden = jnp.eye(2)
    n = c._nqubits
    for i in range(len(arch)):
        u = tc.gates.Gate(unitary[i].reshape(2, 2, 2, 2))
        cand = [tc.gates.Gate(iden), [u[0], u[2]], [u[1], u[3]]]
        mpos = [cand[arch[i, j]] for j in range(n)]
        mpo = tc.quantum.QuOperator([m[0] for m in mpos], [m[1] for m in mpos])
        c.mpo(*range(n), mpo=mpo)
    return c

but it's still non jittable, because cand is not an array. Any thoughts on that? Thanks!

[refraction-ray]

two-qubit gate is more subtle, if the layout is more specified, say brickwall or ladder, i.e. the two qubit gates are always applied in adjacent indices, it is okay to jit using similar approaches. However, if all-to-all couplings are possible to apply two-qubit gates, for now, it is hard to efficiently jit (maybe there is also a clever way, just I fail to recongnize )

[haimengzhao]

I see. Thanks for your suggestions!

[refraction-ray]

btw, i will be at tsinghua U in 10-15 Aug, if you are there by then, we can meet offline. You can drop me an email for wechat id to connect if you are available.