Is there a vectorized variant of the `slice`

function or an alternative to achieve the following task?

Assume that I have a two-dimensional tensor `A`

from which I want to access a slice, e.g.

```
A.index( { Slice(i-1, i+1, 1), Slice(j-1, j+1, 1) } )
```

This works well and gives me a 3x3 sub-tensor with entries `A[i-1:i+1][j-1:j+1]`

as long as `i`

and `j`

are scalars. Instead of looping over a set of `i`

’s and `j`

’s I would like to ‘stack’ all of them along the third dimension so that I can apply some operation (in my case this is a `matmul`

with another three-dimensional tensor) to all of them.

The reason for this request is that one-by-one processing requires me to store the result with `index_put_`

which is leading to problems when calculating gradients (see my other post here).