Definir un gradiente con respecto a un subtensor en Theano

Tengo lo que conceptualmente es una pregunta simple sobre Theano, pero no he podido encontrar la respuesta (confesaré por adelantado que realmente no entiendo cómo funcionan las variables compartidas en Theano, a pesar de muchas horas con los tutoriales).

Estoy tratando de implementar una "red deconvolucional"; específicamente tengo un 3 tensor de entradas (cada entrada es una imagen 2D) y un 4 tensor de códigos; para los iésimos códigos de entrada [i] representa un conjunto de palabras de código que juntas codifican para la entrada i.

He tenido muchos problemas para descubrir cómo hacer un descenso de gradiente en las palabras de código. Aquí están las partes relevantes de mi código:

idx = T.lscalar()
pre_loss_conv = conv2d(input = codes[idx].dimshuffle('x', 0, 1,2),
                       filters = dicts.dimshuffle('x', 0,1, 2),
                       border_mode = 'valid')
loss_conv = pre_loss_conv.reshape((pre_loss_conv.shape[2], pre_loss_conv.shape[3]))
loss_in = inputs[idx]
loss = T.sum(1./2.*(loss_in - loss_conv)**2) 

del_codes = T.grad(loss, codes[idx])
delc_fn = function([idx], del_codes)
train_codes = function([input_index], loss, updates = [
    [codes, T.set_subtensor(codes[input_index], codes[input_index] - 
                            learning_rate*del_codes[input_index])     ]])

(aquí los códigos y los dictos son variables tensoras compartidas). Theano no está contento con esto, específicamente con la definición

del_codes = T.grad(loss, codes[idx])

El mensaje de error que recibo es:theano.gradient.DisconnectedInputError: se le pidió al método graduado que calcule el gradiente con respecto a una variable que no es parte del gráfico computacional del costo, o que solo lo utiliza un operador no diferenciable: Subtensor {int64} .0

Supongo que quiere una variable simbólica en lugar de códigos [idx]; pero entonces no estoy seguro de cómo conectar todo para obtener el efecto deseado. Supongo que tendré que cambiar la línea final a algo así

learning_rate*del_codes)     ]])

¿Alguien puede darme algunos consejos sobre cómo definir esta función correctamente? Creo que probablemente me estoy perdiendo algo básico sobre trabajar con Theano, pero no estoy seguro de qué.

¡Gracias por adelantado!

-Justin

Actualización: la sugerencia de Kyle funcionó muy bien. Aquí está el código específico que usé

current_codes = T.tensor3('current_codes')
current_codes = codes[input_index]
pre_loss_conv = conv2d(input = current_codes.dimshuffle('x', 0, 1,2),
                       filters = dicts.dimshuffle('x', 0,1, 2),
                       border_mode = 'valid')
loss_conv = pre_loss_conv.reshape((pre_loss_conv.shape[2], pre_loss_conv.shape[3]))
loss_in = inputs[input_index]
loss = T.sum(1./2.*(loss_in - loss_conv)**2)  

del_codes = T.grad(loss, current_codes)
train_codes = function([input_index], loss)
train_dicts = theano.function([input_index], loss, updates = [[dicts, dicts - learning_rate*del_dicts]])
codes_update = ( codes, T.set_subtensor(codes[input_index], codes[input_index] - learning_rate*del_codes) )
codes_update_fn = function([input_index], updates = [codes_update])

for i in xrange(num_inputs):
     current_loss = train_codes(i)
     codes_update_fn(i)