Tensor Indexing

The PyTorch C++ API provides tensor indexing similar to Python. Use torch::indexing namespace for index types:

using namespace torch::indexing;

The main difference from Python is that instead of using the [] operator, the C++ API uses the index and index_put_ methods:

• torch::Tensor::index — read elements

• torch::Tensor::index_put_ — write elements

Index Types

The TensorIndex class accepts six types of indices via implicit constructors:

Type	C++	Python equivalent
None (unsqueeze)	None	None
Ellipsis	Ellipsis or "..."	...
Integer	0, 1, -1	0, 1, -1
Boolean	true, false	True, False
Slice	Slice(start, stop, step)	start:stop:step
Tensor	torch::tensor({0, 2})	torch.tensor([0, 2])

Getter Operations

Python	C++
tensor[0]	tensor.index({0})
tensor[-1]	tensor.index({-1})
tensor[1, 2]	tensor.index({1, 2})
tensor[1, :, 3]	tensor.index({1, Slice(), 3})
tensor[None]	tensor.index({None})
tensor[:, None]	tensor.index({Slice(), None})
tensor[...]	tensor.index({Ellipsis}) or tensor.index({"..."})
tensor[..., 0]	tensor.index({Ellipsis, 0})
tensor[1::2]	tensor.index({Slice(1, None, 2)})
tensor[True]	tensor.index({true})
tensor[torch.tensor([1, 2])]	tensor.index({torch::tensor({1, 2})})
tensor[bool_mask]	tensor.index({bool_mask})
tensor[:, torch.tensor([[0,1],[4,3]])]	tensor.index({Slice(), torch::tensor({{0,1},{4,3}})})
tensor[cond > 0]	tensor.index({cond > 0})

Setter Operations

Python	C++
tensor[0] = 1	tensor.index_put_({0}, 1)
tensor[1, 2] = 1	tensor.index_put_({1, 2}, 1)
tensor[1] = torch.arange(5)	tensor.index_put_({1}, torch::arange(5))
tensor[1::2] = 1	tensor.index_put_({Slice(1, None, 2)}, 1)
tensor[0, 1::2] = torch.tensor([3., 4.])	tensor.index_put_({0, Slice(1, None, 2)}, torch::tensor({3., 4.}))
tensor[...] = 0	tensor.index_put_({Ellipsis}, 0)
tensor[None] = value	tensor.index_put_({None}, value)
tensor[bool_mask] = 0	tensor.index_put_({bool_mask}, 0)
tensor[torch.tensor([0, 2])] = value	tensor.index_put_({torch::tensor({0, 2})}, value)
tensor[1:2, torch.tensor([1,2])] = 0	tensor.index_put_({Slice(1, 2), torch::tensor({1, 2})}, 0)

The index_put_ method also accepts an optional accumulate parameter. When true, values are added to existing values instead of replacing them:

tensor.index_put_({mask}, values, /*accumulate=*/true);

Slice Syntax

The Slice constructor signature is:

Slice(
    std::optional<c10::SymInt> start = std::nullopt,
    std::optional<c10::SymInt> stop  = std::nullopt,
    std::optional<c10::SymInt> step  = std::nullopt);

Pass None for open-ended bounds:

Python	C++
: or ::	Slice() or Slice(None, None)
1:	Slice(1, None)
:3	Slice(None, 3)
1:3	Slice(1, 3)
1:3:2	Slice(1, 3, 2)
::2	Slice(None, None, 2)

Full Example

#include <torch/torch.h>

using namespace torch::indexing;

auto tensor = torch::arange(2 * 3 * 4).reshape({2, 3, 4});

// Basic indexing
auto row = tensor.index({0});             // tensor[0]
auto elem = tensor.index({1, 2, 3});      // tensor[1, 2, 3]

// Slicing
auto sliced = tensor.index({Slice(), Slice(0, 2)});  // tensor[:, 0:2]

// None (unsqueeze) and Ellipsis
auto unsqueezed = tensor.index({None});              // tensor[None]
auto last_dim = tensor.index({Ellipsis, -1});         // tensor[..., -1]

// Boolean mask indexing
auto mask = tensor > 10;
auto selected = tensor.index({mask});     // tensor[tensor > 10]

// Integer tensor (fancy) indexing
auto idx = torch::tensor({0, 2});
auto gathered = tensor.index({Slice(), idx});  // tensor[:, [0, 2]]

// Setting values
tensor.index_put_({0, Slice(), 0}, 99);        // tensor[0, :, 0] = 99
tensor.index_put_({mask}, 0);                  // tensor[tensor > 10] = 0