mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 00:21:07 +01:00
36919278cc
Summary:
This PR adds the following items:
- **1st item**: `ArrayRef<TensorIndex>` and `std::initializer_list<TensorIndex>` overloads for `Tensor::index` and `Tensor::index_put_`, to be used specifically for multi-dim indexing purpose.
Design rationale:
* C++ `Tensor::index` and `Tensor::index_put_` are both existing tensor APIs, and they currently (before this PR) only accept a list of tensors (i.e. `ArrayRef<Tensor>`) as indices. If we change their signatures to also accept non-tensors as indices (i.e. `ArrayRef<TensorIndex>`, and `TensorIndex` is convertible from `Tensor` / `Slice` / `None` / `Ellipsis`), it would slow down the original code path (since now it has to go through more steps), which is undesirable.
To get around this problem, the proposed solution is to keep the original `ArrayRef<Tensor>` overload, and add `ArrayRef<TensorIndex>` and `std::initializer_list<TensorIndex>` overloads to `Tensor::index` and `Tensor::index_put_`. This way, the original code path won’t be affected, and the tensor multi-dim indexing API is only used when the user explicitly pass an `ArrayRef<TensorIndex>` or a braced-init-list of `TensorIndex`-convertible types to `Tensor::index` and `Tensor::index_put_` .
Note that the above proposed solution would still affect perf for the user’s original `Tensor::index` or `Tensor::index_put_` call sites that use a braced-init-list of tensors as input, e.g. `tensor.index({...})` or `tensor.index_put_({...}, value)`, since now such function calls would take the multi-dim indexing path instead of the original advanced indexing path. However, there are only two instances of this in our codebase (one in ATen cpp test, one in a C++ API nn init function), and they can be easily changed to explicitly use `ArrayRef<Tensor>` as input (I changed them in this PR). For external user’s code, since this is part of the C++ frontend which is still considered experimental, we will only talk about this change in the release note, and ask users to switch to using `ArrayRef<Tensor>` explicitly if they want to keep using the original advanced indexing code path.
- **2nd item**: Mechanisms for parsing `ArrayRef<TensorIndex>` indices and performing indexing operations (mirroring the functions in `torch/csrc/autograd/python_variable_indexing.cpp`).
- **3rd item**: Simple tests to demonstrate that the `Tensor::index()` and `Tensor::index_put_()` APIs work. I will add more tests after the first few PRs are reviewed.
- **4th item**: Merge Python/C++ indexing code paths, for code simplicity. I tested locally and found that there is no perf regression resulting from the merge. I will get more concrete numbers for common use cases when we settle on the overall design.
This PR supersedes https://github.com/pytorch/pytorch/pull/30425.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/32841
Differential Revision: D19919692
Pulled By: yf225
fbshipit-source-id: 7467e64f97fc0e407624809dd183c95ea16b1482
155 lines
7.1 KiB
C++
155 lines
7.1 KiB
C++
#include <gtest/gtest.h>
|
|
|
|
// TODO: Move the include into `ATen/ATen.h`, once C++ tensor indexing
|
|
// is ready to ship.
|
|
#include <ATen/native/TensorIndexing.h>
|
|
#include <torch/torch.h>
|
|
|
|
#include <test/cpp/api/support.h>
|
|
|
|
using namespace torch::indexing;
|
|
using namespace torch::test;
|
|
|
|
TEST(TensorIndexingTest, Slice) {
|
|
torch::indexing::impl::Slice slice(1, 2, 3);
|
|
ASSERT_EQ(slice.start(), 1);
|
|
ASSERT_EQ(slice.stop(), 2);
|
|
ASSERT_EQ(slice.step(), 3);
|
|
|
|
ASSERT_EQ(c10::str(slice), "1:2:3");
|
|
}
|
|
|
|
TEST(TensorIndexingTest, TensorIndex) {
|
|
{
|
|
std::vector<TensorIndex> indices = {None, "...", Ellipsis, 0, true, {1, None, 2}, torch::tensor({1, 2})};
|
|
ASSERT_TRUE(indices[0].is_none());
|
|
ASSERT_TRUE(indices[1].is_ellipsis());
|
|
ASSERT_TRUE(indices[2].is_ellipsis());
|
|
ASSERT_TRUE(indices[3].is_integer());
|
|
ASSERT_TRUE(indices[3].integer() == 0);
|
|
ASSERT_TRUE(indices[4].is_boolean());
|
|
ASSERT_TRUE(indices[4].boolean() == true);
|
|
ASSERT_TRUE(indices[5].is_slice());
|
|
ASSERT_TRUE(indices[5].slice().start() == 1);
|
|
ASSERT_TRUE(indices[5].slice().stop() == INDEX_MAX);
|
|
ASSERT_TRUE(indices[5].slice().step() == 2);
|
|
ASSERT_TRUE(indices[6].is_tensor());
|
|
ASSERT_TRUE(torch::equal(indices[6].tensor(), torch::tensor({1, 2})));
|
|
}
|
|
|
|
ASSERT_THROWS_WITH(
|
|
TensorIndex(".."),
|
|
"Expected \"...\" to represent an ellipsis index, but got \"..\"");
|
|
|
|
// NOTE: Some compilers such as Clang and MSVC always treat `TensorIndex({1})` the same as
|
|
// `TensorIndex(1)`. This is in violation of the C++ standard
|
|
// (`https://en.cppreference.com/w/cpp/language/list_initialization`), which says:
|
|
// ```
|
|
// copy-list-initialization:
|
|
//
|
|
// U( { arg1, arg2, ... } )
|
|
//
|
|
// functional cast expression or other constructor invocations, where braced-init-list is used
|
|
// in place of a constructor argument. Copy-list-initialization initializes the constructor's parameter
|
|
// (note; the type U in this example is not the type that's being list-initialized; U's constructor's parameter is)
|
|
// ```
|
|
// When we call `TensorIndex({1})`, `TensorIndex`'s constructor's parameter is being list-initialized with {1}.
|
|
// And since we have the `TensorIndex(std::initializer_list<c10::optional<int64_t>>)` constructor, the following
|
|
// rule in the standard applies:
|
|
// ```
|
|
// The effects of list initialization of an object of type T are:
|
|
//
|
|
// if T is a specialization of std::initializer_list, the T object is direct-initialized or copy-initialized,
|
|
// depending on context, from a prvalue of the same type initialized from the braced-init-list.
|
|
// ```
|
|
// Therefore, if the compiler strictly follows the standard, it should treat `TensorIndex({1})` as
|
|
// `TensorIndex(std::initializer_list<c10::optional<int64_t>>({1}))`. However, this is not the case for
|
|
// compilers such as Clang and MSVC, and hence we skip this test for those compilers.
|
|
#if !defined(__clang__) && !defined(_MSC_VER)
|
|
ASSERT_THROWS_WITH(
|
|
TensorIndex({1}),
|
|
"Expected 0 / 2 / 3 elements in the braced-init-list to represent a slice index, but got 1 element(s)");
|
|
#endif
|
|
|
|
ASSERT_THROWS_WITH(
|
|
TensorIndex({1, 2, 3, 4}),
|
|
"Expected 0 / 2 / 3 elements in the braced-init-list to represent a slice index, but got 4 element(s)");
|
|
|
|
{
|
|
std::vector<TensorIndex> indices = {None, "...", Ellipsis, 0, true, {1, None, 2}};
|
|
ASSERT_EQ(c10::str(indices), c10::str("(None, ..., ..., 0, true, 1:", INDEX_MAX, ":2)"));
|
|
ASSERT_EQ(c10::str(indices[0]), "None");
|
|
ASSERT_EQ(c10::str(indices[1]), "...");
|
|
ASSERT_EQ(c10::str(indices[2]), "...");
|
|
ASSERT_EQ(c10::str(indices[3]), "0");
|
|
ASSERT_EQ(c10::str(indices[4]), "true");
|
|
ASSERT_EQ(c10::str(indices[5]), c10::str("1:", INDEX_MAX, ":2"));
|
|
}
|
|
|
|
ASSERT_EQ(c10::str(std::vector<TensorIndex>({{}})), c10::str("(0:", INDEX_MAX, ":1)"));
|
|
ASSERT_EQ(c10::str(std::vector<TensorIndex>({{None, None}})), c10::str("(0:", INDEX_MAX, ":1)"));
|
|
ASSERT_EQ(c10::str(std::vector<TensorIndex>({{None, None, None}})), c10::str("(0:", INDEX_MAX, ":1)"));
|
|
|
|
ASSERT_EQ(c10::str(std::vector<TensorIndex>({{1, None}})), c10::str("(1:", INDEX_MAX, ":1)"));
|
|
ASSERT_EQ(c10::str(std::vector<TensorIndex>({{1, None, None}})), c10::str("(1:", INDEX_MAX, ":1)"));
|
|
ASSERT_EQ(c10::str(std::vector<TensorIndex>({{None, 3}})), c10::str("(0:3:1)"));
|
|
ASSERT_EQ(c10::str(std::vector<TensorIndex>({{None, 3, None}})), c10::str("(0:3:1)"));
|
|
ASSERT_EQ(c10::str(std::vector<TensorIndex>({{None, None, 2}})), c10::str("(0:", INDEX_MAX, ":2)"));
|
|
ASSERT_EQ(c10::str(std::vector<TensorIndex>({{None, None, -1}})), c10::str("(", INDEX_MAX, ":", INDEX_MIN, ":-1)"));
|
|
|
|
ASSERT_EQ(c10::str(std::vector<TensorIndex>({{1, 3}})), c10::str("(1:3:1)"));
|
|
ASSERT_EQ(c10::str(std::vector<TensorIndex>({{1, None, 2}})), c10::str("(1:", INDEX_MAX, ":2)"));
|
|
ASSERT_EQ(c10::str(std::vector<TensorIndex>({{1, None, -1}})), c10::str("(1:", INDEX_MIN, ":-1)"));
|
|
ASSERT_EQ(c10::str(std::vector<TensorIndex>({{None, 3, 2}})), c10::str("(0:3:2)"));
|
|
ASSERT_EQ(c10::str(std::vector<TensorIndex>({{None, 3, -1}})), c10::str("(", INDEX_MAX, ":3:-1)"));
|
|
|
|
ASSERT_EQ(c10::str(std::vector<TensorIndex>({{1, 3, 2}})), c10::str("(1:3:2)"));
|
|
}
|
|
|
|
TEST(TensorIndexingTest, TestAdvancedIndexingWithArrayRefOfTensor) {
|
|
{
|
|
torch::Tensor tensor = torch::randn({20, 20});
|
|
torch::Tensor index = torch::arange(10, torch::kLong).cpu();
|
|
torch::Tensor result_with_array_ref = tensor.index(at::ArrayRef<torch::Tensor>({index}));
|
|
torch::Tensor result_with_init_list = tensor.index({index});
|
|
ASSERT_TRUE(result_with_array_ref.equal(result_with_init_list));
|
|
}
|
|
{
|
|
torch::Tensor tensor = torch::randn({20, 20});
|
|
torch::Tensor index = torch::arange(10, torch::kLong).cpu();
|
|
torch::Tensor result_with_array_ref = tensor.index_put_(at::ArrayRef<torch::Tensor>({index}), torch::ones({20}));
|
|
torch::Tensor result_with_init_list = tensor.index_put_({index}, torch::ones({20}));
|
|
ASSERT_TRUE(result_with_array_ref.equal(result_with_init_list));
|
|
}
|
|
{
|
|
torch::Tensor tensor = torch::randn({20, 20});
|
|
torch::Tensor index = torch::arange(10, torch::kLong).cpu();
|
|
torch::Tensor result_with_array_ref = tensor.index_put_(at::ArrayRef<torch::Tensor>({index}), torch::ones({1, 20}));
|
|
torch::Tensor result_with_init_list = tensor.index_put_({index}, torch::ones({1, 20}));
|
|
ASSERT_TRUE(result_with_array_ref.equal(result_with_init_list));
|
|
}
|
|
}
|
|
|
|
TEST(TensorIndexingTest, TestSingleInt) {
|
|
auto v = torch::randn({5, 7, 3});
|
|
ASSERT_EQ(v.index({4}).sizes(), torch::IntArrayRef({7, 3}));
|
|
}
|
|
|
|
TEST(TensorIndexingTest, TestMultipleInt) {
|
|
auto v = torch::randn({5, 7, 3});
|
|
ASSERT_EQ(v.index({4}).sizes(), torch::IntArrayRef({7, 3}));
|
|
ASSERT_EQ(v.index({4, {}, 1}).sizes(), torch::IntArrayRef({7}));
|
|
|
|
// To show that `.index_put_` works
|
|
v.index_put_({4, 3, 1}, 0);
|
|
ASSERT_EQ(v.index({4, 3, 1}).item<double>(), 0);
|
|
}
|
|
|
|
TEST(TensorIndexingTest, TestNone) {
|
|
auto v = torch::randn({5, 7, 3});
|
|
ASSERT_EQ(v.index({None}).sizes(), torch::IntArrayRef({1, 5, 7, 3}));
|
|
ASSERT_EQ(v.index({{}, None}).sizes(), torch::IntArrayRef({5, 1, 7, 3}));
|
|
ASSERT_EQ(v.index({{}, None, None}).sizes(), torch::IntArrayRef({5, 1, 1, 7, 3}));
|
|
ASSERT_EQ(v.index({"...", None}).sizes(), torch::IntArrayRef({5, 7, 3, 1}));
|
|
}
|