Differentiable Indexing

[magz]

I think you are looking for this? Indexing a variable with a variable