PyTorch C++ extension: How to index tensor and update it?

Question

I'm creating a PyTorch C++ extension and after much research I can't figure out how to index a tensor and update its values. I found out how to iterate over a tensor's entries using the data_ptr() method, but that's not applicable to my use case.

Given is a matrix M, a list of lists (blocks) of index pairs P and a function f: dtype(M)^2 -> dtype(M)^2 that takes two values and spits out two new values.

I'm trying to implement the following pseudo code:

for each block B in P:
    for each row R in M:
        for each index-pair (i,j) in B:
            M[R,i], M[R,j] = f(M[R,i], M[R,j])

After all, this code is going to run on the GPU using CUDA, but since I don't have any experience with that, I wanted to first write a pure C++ program and then convert it.

Can anyone suggest how to do this or how to convert the algorithm to do something equivalent?

Niki Niki · Accepted Answer · 2020-09-08T07:54:11

What I wanted to do can be done using the tensor.accessor<scalar_dtype, num_dimensions>() method. If executing on the GPU instead use scalars.packed_accessor64<scalar_dtype, num_dimensions, torch::RestrictPtrTraits>() or scalars.packed_accessor32<scalar_dtype, num_dimensions, torch::RestrictPtrTraits>() (depending on the size of your tensor).

auto num_rows = scalars.size(0);
matrix = torch::rand({10, 8});
auto a = matrix.accessor<float, 2>();
for (auto i = 0; i < num_rows; ++i) {
    auto x = a[i][some_index];
    auto new_x = some_function(x);
    a[i][some_index] = new_x;
}

PyTorch C++ extension: How to index tensor and update it?

1 Answers