Deterministic Theano Op - Gradient

Thank you for your suggestion!
Could you give me a hint as to how I could create such an index matrix?
Since Metropolis converges slowly on my dataset, I’d like to use NUTS, but I find it difficult to get started with Theano.