Definindo um gradiente com relação a um subtensor no Theano

Eu tenho o que é conceitualmente uma pergunta simples sobre o Theano, mas não consegui encontrar a resposta (confesso que não entendo realmente como as variáveis compartilhadas funcionam no Theano, apesar de muitas horas com os tutoriais).

Estou tentando implementar uma "rede deconvolucional"; especificamente, eu tenho um tensor de 3 entradas (cada entrada é uma imagem 2D) e um tensor de 4 códigos; para os enésimos códigos de entrada [i] representa um conjunto de palavras de código que juntas codificam para a entrada i.

Tenho tido muitos problemas para descobrir como fazer a descida de gradiente nas palavras de código. Aqui estão as partes relevantes do meu 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])     ]])

(aqui códigos e dictos são variáveis compartilhadas do tensor). Theano está descontente com isso, especificamente ao definir

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

A mensagem de erro que estou recebendo é:O método theano.gradient.DisconnectedInputError: grad foi solicitado a calcular o gradiente em relação a uma variável que não faz parte do gráfico computacional do custo ou é usada apenas por um operador não diferenciável: Subtensor {int64} .0

Eu estou supondo que ele quer uma variável simbólica em vez de códigos [idx]; mas não sei como conectar tudo para obter o efeito pretendido. Acho que vou precisar mudar a linha final para algo como

learning_rate*del_codes)     ]])

Alguém pode me dar algumas dicas sobre como definir essa função corretamente? Eu acho que provavelmente estou perdendo algo básico sobre trabalhar com Theano, mas não sei ao certo o que.

Desde já, obrigado!

-Justin

Atualização: A sugestão de Kyle funcionou muito bem. Aqui está o código específico que eu usei

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)

questionAnswers(1)

yourAnswerToTheQuestion