Problem updating tensor: x must be the result of a subtensor operation


I am trying to implement for pymc3 function which uses products of complex tensors to describe the propagation of a wave in a stratified medium. I have a working implementation in numpy/cython but want to use the full Theano potential.

Here is a stripped down version of my code, which raises a TypeError:x must be the result of a subtensor operation in the last block

q = T.scalar("q")
kz = q / 2

sld_sub = 3e-6
sld = p.array([0, 1e-6, sld_sub]) + 0j
thick = p.array([0, 100])

kn = T.sqrt(kz**2 - 4 * p.pi * (sld - sld[0]))

d = theano.tensor.extra_ops.diff(kn)
s = kn[:-1] + kn[1:]
r_n_np = d / s

beta = kn[:-1] * thick * 1j

expb = theano.tensor.exp(beta)
exp_b = 1. / expb

M = p.zeros((len(thick),2, 2), dtype=theano.tensor.ztensor3)

M = T.set_subtensor(M[:,0,0], expb)
M = T.set_subtensor(M[:,0,1], r_n_np * expb)
M = T.set_subtensor(M[:,1,0], r_n_np * exp_b)    
M = T.set_subtensor(M[:,1,1], exp_b)

Once this error is lifted, I then intend to perform the equivalent of

prod_M = np.linalg.multidot(M)

I guess I will have to use scan to achieve that.

Eventually I need then to perform some reduce function leading to

R = (abs(prod_M[1,0] / prod_M[0,0]))**2

The complete code above should work for a scalar value of q I would then use scan to perform f(q=vector).

Thanks in advance for helping me solve the TypeError problem, and also for any tip regarding the next steps!