Advanced Tensor slicing in C++

I’m trying to figure out how to do the following in the C++ Frontend

I have an (21,3) tensor

xyz=torch.tensor(
[[-31.986  69.746 353.42 ]
 [-37.302  31.703 339.31 ]
 [-56.901   5.73  330.64 ]
 [-54.698 -23.287 323.1  ]
 [-38.11  -42.968 316.82 ]
 [-18.578 -11.142 369.19 ]
 [ -0.725 -41.328 361.34 ]
 [ 12.068 -60.113 353.56 ]
 [ 24.138 -78.254 350.26 ]
 [ -0.     -0.    375.   ]
 [ 22.376 -33.838 388.76 ]
 [ 36.284 -56.047 381.05 ]
 [ 46.965 -73.672 368.   ]
 [ 14.386  10.085 375.35 ]
 [ 36.256 -22.517 385.73 ]
 [ 49.229 -42.654 382.12 ]
 [ 50.176 -49.916 360.85 ]
 [ 25.597  23.502 372.95 ]
 [ 48.176  11.614 364.04 ]
 [ 65.111   2.527 357.42 ]
 [ 83.033  -6.355 354.3  ]])

I want to convert it into a (5,5,3) array like below

[[[-31.986  69.746 353.42 ]
  [-37.302  31.703 339.31 ]
  [-56.901   5.73  330.64 ]
  [-54.698 -23.287 323.1  ]
  [-38.11  -42.968 316.82 ]]

 [[-31.986  69.746 353.42 ]
  [-18.578 -11.142 369.19 ]
  [ -0.725 -41.328 361.34 ]
  [ 12.068 -60.113 353.56 ]
  [ 24.138 -78.254 350.26 ]]

 [[-31.986  69.746 353.42 ]
  [ -0.     -0.    375.   ]
  [ 22.376 -33.838 388.76 ]
  [ 36.284 -56.047 381.05 ]
  [ 46.965 -73.672 368.   ]]

 [[-31.986  69.746 353.42 ]
  [ 14.386  10.085 375.35 ]
  [ 36.256 -22.517 385.73 ]
  [ 49.229 -42.654 382.12 ]
  [ 50.176 -49.916 360.85 ]]

 [[-31.986  69.746 353.42 ]
  [ 25.597  23.502 372.95 ]
  [ 48.176  11.614 364.04 ]
  [ 65.111   2.527 357.42 ]
  [ 83.033  -6.355 354.3  ]]]

Normally I would do this in python with advance slicing using forexample

indexing = torch.tensor([[0,1 ,2 ,3 ,4] ,  
                         [0,5 ,6 ,7 ,8] , 
                         [0,9 ,10,11,12],  
                         [0,13,14,15,16],  
                         [0,17,18,19,20]]) 

xyz[indexing]

How would I do something similar in C++? I’ve seen that the function `.slice(/dim/, /start/, /stop/, /step/) , but I’m not sure how I would acheive what I want using this

Thank you for the help!

The C++ equivalent of [...] indexing is torch::index(self, indices) where indices is a vector of Tensors (which can be undefined).
In your case:

std::vector<torch::Tensor> idx_v{indexing};
auto result = torch::index(xyz, idx_v);

Best regards

Thomas

2 Likes

Worked perfectly, thank you! Incase anyone else reads this make sure that you’re indexing tensor is of type torch::kInt64

1 Like