Wrap around negative indices in the index method

Author: asuhan

test/cpp/test_aten_xla_tensor.cpp

@@ -1259,7 +1259,7 @@ TEST_F(AtenXlaTensorTest, TestBroadcastTensors) {
 TEST_F(AtenXlaTensorTest, TestOneIndex) {
   at::Tensor params = at::rand({4, 3, 5, 6, 7}, at::TensorOptions(at::kFloat));
   at::Tensor indices =
-      at::randint(0, 3, {2, 4, 3}, at::TensorOptions(at::kLong));
+      at::randint(-3, 3, {2, 4, 3}, at::TensorOptions(at::kLong));
   at::Tensor result = at::index(params, {indices});
   ForEachDevice([&](const Device& device) {
     at::Tensor xla_params = bridge::CreateXlaTensor(params, device);
@@ -1272,10 +1272,10 @@ TEST_F(AtenXlaTensorTest, TestOneIndex) {
 TEST_F(AtenXlaTensorTest, TestMultiIndexMiddleNull) {
   at::Tensor params = at::rand({4, 3, 5, 6, 7}, at::TensorOptions(at::kFloat));
   at::Tensor indices_0 =
-      at::randint(0, 3, {2, 4, 3}, at::TensorOptions(at::kLong));
+      at::randint(-3, 3, {2, 4, 3}, at::TensorOptions(at::kLong));
   at::Tensor indices_null;
   at::Tensor indices_1 =
-      at::randint(0, 3, {2, 4, 3}, at::TensorOptions(at::kLong));
+      at::randint(-3, 3, {2, 4, 3}, at::TensorOptions(at::kLong));
   at::Tensor result = at::index(params, {indices_0, indices_null, indices_1});
   ForEachDevice([&](const Device& device) {
     at::Tensor xla_params = bridge::CreateXlaTensor(params, device);
@@ -1290,10 +1290,10 @@ TEST_F(AtenXlaTensorTest, TestMultiIndexMiddleNull) {
 TEST_F(AtenXlaTensorTest, TestMultiIndexTailNull) {
   at::Tensor params = at::rand({4, 3, 5, 6, 7}, at::TensorOptions(at::kFloat));
   at::Tensor indices_0 =
-      at::randint(0, 3, {2, 4, 3}, at::TensorOptions(at::kLong));
+      at::randint(-3, 3, {2, 4, 3}, at::TensorOptions(at::kLong));
   at::Tensor indices_null;
   at::Tensor indices_1 =
-      at::randint(0, 3, {2, 4, 3}, at::TensorOptions(at::kLong));
+      at::randint(-3, 3, {2, 4, 3}, at::TensorOptions(at::kLong));
   at::Tensor result = at::index(params, {indices_0, indices_1, indices_null});
   ForEachDevice([&](const Device& device) {
     at::Tensor xla_params = bridge::CreateXlaTensor(params, device);
@@ -1308,9 +1308,9 @@ TEST_F(AtenXlaTensorTest, TestMultiIndexTailNull) {
 TEST_F(AtenXlaTensorTest, TestMultiIndexMiddleBroadcast) {
   at::Tensor params = at::rand({4, 3, 5, 6, 7}, at::TensorOptions(at::kFloat));
   at::Tensor indices_0 =
-      at::randint(0, 3, {2, 4, 3}, at::TensorOptions(at::kLong));
+      at::randint(-3, 3, {2, 4, 3}, at::TensorOptions(at::kLong));
   at::Tensor indices_1 =
-      at::randint(0, 3, {2, 1, 3}, at::TensorOptions(at::kLong));
+      at::randint(-3, 3, {2, 1, 3}, at::TensorOptions(at::kLong));
   at::Tensor result = at::index(params, {indices_0, indices_1});
   ForEachDevice([&](const Device& device) {
     at::Tensor xla_params = bridge::CreateXlaTensor(params, device);
@@ -1325,9 +1325,9 @@ TEST_F(AtenXlaTensorTest, TestMultiIndexMiddleBroadcast) {
 TEST_F(AtenXlaTensorTest, TestMultiIndexTailBroadcast) {
   at::Tensor params = at::rand({4, 3, 5, 6, 7}, at::TensorOptions(at::kFloat));
   at::Tensor indices_0 =
-      at::randint(0, 3, {2, 1, 3}, at::TensorOptions(at::kLong));
+      at::randint(-3, 3, {2, 1, 3}, at::TensorOptions(at::kLong));
   at::Tensor indices_1 =
-      at::randint(0, 3, {2, 1}, at::TensorOptions(at::kLong));
+      at::randint(-3, 3, {2, 1}, at::TensorOptions(at::kLong));
   at::Tensor result = at::index(params, {indices_0, indices_1});
   ForEachDevice([&](const Device& device) {
     at::Tensor xla_params = bridge::CreateXlaTensor(params, device);

torch_xla/csrc/tensor.cpp

@@ -1045,10 +1045,11 @@ XLATensor XLATensor::expand(const XLATensor& input,
 XLATensor XLATensor::index(
     const XLATensor& input,
     tensorflow::gtl::ArraySlice<const XLATensor> indices) {
-  xla::int64 indices_rank = indices.front().shape().get().rank();
+  auto canonical_indices = WrapIndicesOnce(input, indices);
+  xla::int64 indices_rank = canonical_indices.front().shape().get().rank();
   // Stack the indices to allow the whole multi-indexing to be dispatched with a
   // single gather.
-  XLATensor indices_nd = XLATensor::stack(indices, indices_rank);
+  XLATensor indices_nd = XLATensor::stack(canonical_indices, indices_rank);
   return Create(ir::ops::IndexOp(input.GetIrValue(), indices_nd.GetIrValue()),
                 input.GetDevice());
 }
@@ -1869,6 +1870,24 @@ Device XLATensor::CommonDeviceForTensors(
   return device;
 }
 
+std::vector<XLATensor> XLATensor::WrapIndicesOnce(
+    const XLATensor& input,
+    tensorflow::gtl::ArraySlice<const XLATensor> indices) {
+  std::vector<XLATensor> canonical_indices;
+  XLA_CHECK_LE(indices.size(), input.shape().get().rank());
+  for (size_t dim_idx = 0; dim_idx < indices.size(); ++dim_idx) {
+    const XLATensor& dim_index = indices[dim_idx];
+    int64_t dim_size = input.shape().get().dimensions(dim_idx);
+    XLATensor wrapped_dim_index =
+        Create(dim_index.GetIrValue() +
+                   ir::ops::ScalarOp(at::Scalar(dim_size), dim_index.shape()),
+               input.GetDevice());
+    XLATensor wrap_cond = lt(indices[dim_idx], at::Scalar(int64_t(0)));
+    canonical_indices.push_back(where(wrap_cond, wrapped_dim_index, dim_index));
+  }
+  return canonical_indices;
+}
+
 xla::int64 XLATensor::GetCanonicalDimension(const XLATensor& input,
                                             xla::int64 dim) {
   return XlaHelpers::GetCanonicalDimensionIndex(dim,

torch_xla/csrc/tensor.h

@@ -697,6 +697,12 @@ class XLATensor {
 
   static xla::int64 GetNextTensorId();
 
+  // Wraps index tensors once into the [0, dim_size) interval, where dim_size is
+  // the size of the current indexed dimension.
+  static std::vector<XLATensor> WrapIndicesOnce(
+      const XLATensor& input,
+      tensorflow::gtl::ArraySlice<const XLATensor> indices);
+
   static xla::int64 GetCanonicalDimension(const XLATensor& input,
                                           xla::int64 dim);