lasagne - Theano: pass time step argument to step function in custom GRU -

i new theano , lasagne , working on project requires implementing custom gru.

i having problems passing index of scan function step function in code below.

i have based code on this theano-users thread, new theano , missing something.

also, looked @ stackoverflow's theano-pass-iterating-index-of-scan-to-the-called-function, should identical trying do, unable accepted solution work code.

the relevant code snippets below.

i've included 2 step functions below because conditionals determine if step or step_masked passed fn argument theano.scan. i.e. need add time step idx argument both of them.)

please let me know if need provide other context.

class customgrumodel(grulayer):     ...     def get_output_for():         ...          def step(idx, input_n, hid_prev, *args):             # conditionals involving idx             ...          def step_masked(idx, input_n, mask_n, hid_prev, *args):             hid = step(idx, input_n, hid_previous, *args)             ...          # scan op iterated on first dimension of input , repeatedly applies step function         idx = t.scalar('idx', dtype='int64')         sequences.insert(0, t.arange(idx.shape[0]))  # how pass theano.tensor.arange 1 of sequences scan?         non_seqs.insert(0, idx)         hid_out = theano.scan(             fn=step,             sequences=sequences,             go_backwards=self.backwards,             non_sequences = non_seqs,             truncate_gradient=self.gradient_steps,             strict=true,)[0]         ...  class main():     ...     idx = t.scalar('idx', dtype='int64')  # should create shared variable here instead?     train_fn = theano.function([idx, in_x1, in_mask1, in_x2, in_mask2, in_x3, in_mask3, in_y],                                          loss, updates=updates, on_unused_input='warn')