Support negative indices and omitted unit dimensions in index_put_ (#23529)

### Summary

Transformers Mistral sliding-window decode fails after reinplace
rewrites its cache updates to `aten::index_put_`. The portable kernel
rejects the valid index `-1` and requires values `[1, 2, 1, 8]` where
the exported graph supplies `[2, 1, 8]`.

Accept indices in `[-size, size)` and normalize negative indices in both
overwrite and accumulation paths. Allow omitted leading unit dimensions
in the values tensor, using the index length to drive the existing copy
loops. This requires no additional allocation or data expansion.

Depends on #23506 (`JakeStevens:fix/accumulating-index-put-inplace`).
Review that first; the follow-up changes are in commit `95badb36fe`.

### Test plan

Compiled and ran the complete `op_index_put_test.cpp` against the
portable kernel with a local CMake harness
(`./build/fixed_index_put_test`): **57 passed, 2 existing dynamic-shape
skips**. Added 33 regression cases covering int32/int64 indices,
overwrite and accumulation, repeated index aliases, bounds, omitted unit
dimensions, scalar updates, and incompatible shapes.

The corresponding repository tests can be run from a CMake build
configured with `EXECUTORCH_BUILD_TESTS=ON`:

```bash
cmake --build cmake-out --target portable_kernels_test -j
cmake-out/kernels/test/portable_kernels_test --gtest_filter='*OpIndexPut*'
```

Relinked the local runtime with the changed kernel and ran tiny FP32
Mistral static sliding-window single-token decode exports generated with
the planner fix from #23526. All **11 runtime outputs match default
bit-for-bit**, including caches and repeated calls after restoring
inputs. Both paths match eager PyTorch with maximum absolute error
**5.96e-8**.

`lintrunner kernels/portable/cpu/op_index_put.cpp
kernels/test/op_index_put_test.cpp` and the pre-commit checks pass.

Authored with assistance from OpenAI Codex.

Author: JakeStevens

kernels/portable/cpu/op_index_put.cpp

@@ -200,7 +200,7 @@ bool check_special_case_in_place_args(
       Long, Int, index.scalar_type(), ctx, "index_put_", CTYPE, [&]() {
         const CTYPE* const index_arr = index.const_data_ptr<CTYPE>();
         for (const auto i : c10::irange(index.numel())) {
-          if (index_arr[i] < 0 ||
+          if (index_arr[i] < -static_cast<CTYPE>(in.size(*dim)) ||
               index_arr[i] >= static_cast<CTYPE>(in.size(*dim))) {
             ET_LOG(
                 Error,
@@ -220,14 +220,6 @@ bool check_special_case_in_place_args(
       is_valid_index,
       "Some index values are not within bounds of input tensor at indexed dim");
 
-  ET_CHECK_OR_RETURN_FALSE(
-      values.size(*dim) == index.size(0),
-      "Special case in-place index_put requires values to match index length at the indexed dim; values.size(%zu) = %" ET_PRI_TENSOR_SIZE
-      ", index_length = %zd",
-      *dim,
-      values.size(*dim),
-      index.size(0));
-
   Tensor::SizesType expected_values_size[kTensorDimensionLimit] = {};
   size_t in_ndim = static_cast<size_t>(in.dim());
   for (const auto i : c10::irange(in_ndim)) {
@@ -237,6 +229,15 @@ bool check_special_case_in_place_args(
   }
   expected_values_size[*dim] = static_cast<Tensor::SizesType>(index.size(0));
 
+  size_t leading_unit_dims = 0;
+  while (in_ndim - leading_unit_dims > static_cast<size_t>(values.dim()) &&
+         expected_values_size[leading_unit_dims] == 1) {
+    ++leading_unit_dims;
+  }
+  const bool values_shape_matches = tensor_has_expected_size(
+      values,
+      {expected_values_size + leading_unit_dims, in_ndim - leading_unit_dims});
+
 #if ET_LOG_ENABLED
   auto in_shape_str = executorch::runtime::tensor_shape_to_c_string(
       executorch::runtime::Span<const Tensor::SizesType>(
@@ -246,14 +247,14 @@ bool check_special_case_in_place_args(
           values.sizes().data(), values.sizes().size()));
 
   ET_CHECK_OR_RETURN_FALSE(
-      tensor_has_expected_size(values, {expected_values_size, in_ndim}),
-      "Special case in-place index_put requires values to match input shape except for indexed dim; got input shape %s and values shape %s",
+      values_shape_matches,
+      "Special case in-place index_put requires values to match input shape except for indexed dim, allowing omitted leading unit dimensions; got input shape %s and values shape %s",
       in_shape_str.data(),
       values_shape_str.data());
 #else
   ET_CHECK_OR_RETURN_FALSE(
-      tensor_has_expected_size(values, {expected_values_size, in_ndim}),
-      "Special case in-place index_put requires values to match input shape except for indexed dim");
+      values_shape_matches,
+      "Special case in-place index_put requires values to match input shape except for indexed dim, allowing omitted leading unit dimensions");
 #endif // ET_LOG_ENABLED
 
   return true;
@@ -299,7 +300,7 @@ Tensor& index_put_(
     return in;
   }
 
-  size_t values_dim_length = values.size(dim);
+  size_t values_dim_length = index.size(0);
   size_t in_dim_length = in.size(dim);
 
   size_t length_per_step = trailing_dims * in.element_size();
@@ -316,10 +317,13 @@ Tensor& index_put_(
             VALUE_T* dest = in.mutable_data_ptr<VALUE_T>();
             for (const auto i : c10::irange(leading_dims)) {
               for (const auto j : c10::irange(values_dim_length)) {
+                const auto ix = index_arr[j] < 0
+                    ? index_arr[j] + static_cast<CTYPE>(in_dim_length)
+                    : index_arr[j];
                 const size_t src_offset =
                     (i * values_dim_length + j) * trailing_dims;
                 const size_t dest_offset =
-                    (i * in_dim_length + index_arr[j]) * trailing_dims;
+                    (i * in_dim_length + ix) * trailing_dims;
                 for (const auto k : c10::irange(trailing_dims)) {
                   dest[dest_offset + k] += src[src_offset + k];
                 }
@@ -331,8 +335,11 @@ Tensor& index_put_(
         const char* src = values_data + i * values_dim_length * length_per_step;
         char* dest = in_data + i * in_dim_length * length_per_step;
         for (const auto j : c10::irange(values_dim_length)) {
+          const auto ix = index_arr[j] < 0
+              ? index_arr[j] + static_cast<CTYPE>(in_dim_length)
+              : index_arr[j];
           const char* copy_src = src + j * length_per_step;
-          char* copy_dest = dest + index_arr[j] * length_per_step;
+          char* copy_dest = dest + ix * length_per_step;
           memcpy(copy_dest, copy_src, length_per_step);
         }
       }

kernels/test/op_index_put_test.cpp

@@ -16,6 +16,8 @@
 
 #include <gtest/gtest.h>
 #include <sys/types.h>
+#include <tuple>
+#include <vector>
 
 using namespace ::testing;
 using executorch::aten::ArrayRef;
@@ -1138,3 +1140,179 @@ TEST_F(OpIndexPutInplaceTest, AccumulateHalfRepeatedIndices) {
   op_index_put_(x, indices, values, /*accumulate=*/true);
   EXPECT_TENSOR_EQ(x, tf.make({3}, {1, 6, 1}));
 }
+
+class OpIndexPutInplaceNegativeIndexTest
+    : public OpIndexPutInplaceTest,
+      public WithParamInterface<std::tuple<ScalarType, bool>> {};
+
+TEST_P(OpIndexPutInplaceNegativeIndexTest, MixedIndicesInMiddleDimension) {
+  const auto [index_type, accumulate] = GetParam();
+  TensorFactory<ScalarType::Float> tf;
+  TensorFactory<ScalarType::Long> tfl;
+  TensorFactory<ScalarType::Int> tfi;
+  Tensor index = index_type == ScalarType::Long ? tfl.make({3}, {-1, -3, 1})
+                                                : tfi.make({3}, {-1, -3, 1});
+  optional<Tensor> indices[] = {std::nullopt, index};
+  Tensor x = tf.ones({2, 3, 2});
+  Tensor values = tf.make({2, 3, 2}, {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12});
+  Tensor expected = accumulate
+      ? tf.make({2, 3, 2}, {4, 5, 6, 7, 2, 3, 10, 11, 12, 13, 8, 9})
+      : tf.make({2, 3, 2}, {3, 4, 5, 6, 1, 2, 9, 10, 11, 12, 7, 8});
+
+  Tensor ret = op_index_put_(x, indices, values, accumulate);
+  EXPECT_TENSOR_EQ(ret, x);
+  EXPECT_TENSOR_EQ(x, expected);
+  EXPECT_TENSOR_EQ(
+      index,
+      index_type == ScalarType::Long ? tfl.make({3}, {-1, -3, 1})
+                                     : tfi.make({3}, {-1, -3, 1}));
+}
+
+TEST_P(OpIndexPutInplaceNegativeIndexTest, NegativeIndexOutOfBounds) {
+  const auto [index_type, accumulate] = GetParam();
+  TensorFactory<ScalarType::Float> tf;
+  TensorFactory<ScalarType::Long> tfl;
+  TensorFactory<ScalarType::Int> tfi;
+  Tensor index = index_type == ScalarType::Long ? tfl.make({2}, {-1, -4})
+                                                : tfi.make({2}, {-1, -4});
+  optional<Tensor> indices[] = {std::nullopt, index};
+  Tensor x = tf.ones({2, 3, 2});
+  Tensor values = tf.zeros({2, 2, 2});
+
+  ET_EXPECT_KERNEL_FAILURE(
+      context_, op_index_put_(x, indices, values, accumulate));
+#ifndef USE_ATEN_LIB
+  EXPECT_TENSOR_EQ(x, tf.ones({2, 3, 2}));
+#endif
+}
+
+TEST_P(OpIndexPutInplaceNegativeIndexTest, PositiveIndexOutOfBounds) {
+  const auto [index_type, accumulate] = GetParam();
+  TensorFactory<ScalarType::Float> tf;
+  TensorFactory<ScalarType::Long> tfl;
+  TensorFactory<ScalarType::Int> tfi;
+  Tensor index = index_type == ScalarType::Long ? tfl.make({2}, {-1, 3})
+                                                : tfi.make({2}, {-1, 3});
+  optional<Tensor> indices[] = {std::nullopt, index};
+  Tensor x = tf.ones({2, 3, 2});
+  Tensor values = tf.zeros({2, 2, 2});
+
+  ET_EXPECT_KERNEL_FAILURE(
+      context_, op_index_put_(x, indices, values, accumulate));
+#ifndef USE_ATEN_LIB
+  EXPECT_TENSOR_EQ(x, tf.ones({2, 3, 2}));
+#endif
+}
+
+TEST_P(OpIndexPutInplaceNegativeIndexTest, NegativeIndexIntoEmptyDimension) {
+  const auto [index_type, accumulate] = GetParam();
+  TensorFactory<ScalarType::Float> tf;
+  TensorFactory<ScalarType::Long> tfl;
+  TensorFactory<ScalarType::Int> tfi;
+  Tensor index = index_type == ScalarType::Long ? tfl.make({1}, {-1})
+                                                : tfi.make({1}, {-1});
+  Tensor x = tf.zeros({0});
+  Tensor values = tf.ones({1});
+
+  ET_EXPECT_KERNEL_FAILURE(
+      context_, op_index_put_(x, {index}, values, accumulate));
+}
+
+TEST_P(OpIndexPutInplaceNegativeIndexTest, OmittedLeadingUnitDimensions) {
+  const auto [index_type, accumulate] = GetParam();
+  TensorFactory<ScalarType::Float> tf;
+  TensorFactory<ScalarType::Long> tfl;
+  TensorFactory<ScalarType::Int> tfi;
+  Tensor index = index_type == ScalarType::Long ? tfl.make({1}, {-1})
+                                                : tfi.make({1}, {-1});
+  Tensor values = tf.make({2, 1, 2}, {2, 3, 4, 5});
+
+  for (const std::vector<int32_t>& shape :
+       {std::vector<int32_t>{1, 2, 4, 2}, {1, 1, 2, 4, 2}}) {
+    Tensor x = tf.ones(shape);
+    std::vector<optional<Tensor>> indices(shape.size() - 1, std::nullopt);
+    indices.back() = index;
+    Tensor expected = accumulate
+        ? tf.make(shape, {1, 1, 1, 1, 1, 1, 3, 4, 1, 1, 1, 1, 1, 1, 5, 6})
+        : tf.make(shape, {1, 1, 1, 1, 1, 1, 2, 3, 1, 1, 1, 1, 1, 1, 4, 5});
+
+    Tensor ret =
+        op_index_put_(x, {indices.data(), indices.size()}, values, accumulate);
+    EXPECT_TENSOR_EQ(ret, x);
+    EXPECT_TENSOR_EQ(x, expected);
+  }
+  EXPECT_TENSOR_EQ(values, tf.make({2, 1, 2}, {2, 3, 4, 5}));
+}
+
+TEST_P(OpIndexPutInplaceNegativeIndexTest, OmittedIndexedLeadingDimension) {
+  const auto [index_type, accumulate] = GetParam();
+  TensorFactory<ScalarType::Float> tf;
+  TensorFactory<ScalarType::Long> tfl;
+  TensorFactory<ScalarType::Int> tfi;
+  Tensor index = index_type == ScalarType::Long ? tfl.make({1}, {-1})
+                                                : tfi.make({1}, {-1});
+  Tensor x = tf.ones({3, 1, 2});
+  Tensor values = tf.make({2}, {2, 3});
+  Tensor expected = accumulate ? tf.make({3, 1, 2}, {1, 1, 1, 1, 3, 4})
+                               : tf.make({3, 1, 2}, {1, 1, 1, 1, 2, 3});
+
+  op_index_put_(x, {index}, values, accumulate);
+  EXPECT_TENSOR_EQ(x, expected);
+}
+
+TEST_P(OpIndexPutInplaceNegativeIndexTest, ScalarValuesForSingleElement) {
+  const auto [index_type, accumulate] = GetParam();
+  TensorFactory<ScalarType::Float> tf;
+  TensorFactory<ScalarType::Long> tfl;
+  TensorFactory<ScalarType::Int> tfi;
+  Tensor index = index_type == ScalarType::Long ? tfl.make({1}, {-1})
+                                                : tfi.make({1}, {-1});
+  optional<Tensor> indices[] = {std::nullopt, index};
+  Tensor x = tf.ones({1, 3, 1});
+  Tensor values = tf.make({}, {5});
+  Tensor expected = accumulate ? tf.make({1, 3, 1}, {1, 1, 6})
+                               : tf.make({1, 3, 1}, {1, 1, 5});
+
+  op_index_put_(x, indices, values, accumulate);
+  EXPECT_TENSOR_EQ(x, expected);
+}
+
+TEST_P(OpIndexPutInplaceNegativeIndexTest, IncompatibleValuesShape) {
+  const auto [index_type, accumulate] = GetParam();
+  TensorFactory<ScalarType::Float> tf;
+  TensorFactory<ScalarType::Long> tfl;
+  TensorFactory<ScalarType::Int> tfi;
+  Tensor index = index_type == ScalarType::Long ? tfl.make({1}, {-1})
+                                                : tfi.make({1}, {-1});
+  optional<Tensor> indices[] = {std::nullopt, std::nullopt, index};
+  Tensor x = tf.ones({1, 2, 4, 2});
+  Tensor values = tf.zeros({2, 2, 1});
+
+  ET_EXPECT_KERNEL_FAILURE(
+      context_, op_index_put_(x, indices, values, accumulate));
+#ifndef USE_ATEN_LIB
+  EXPECT_TENSOR_EQ(x, tf.ones({1, 2, 4, 2}));
+#endif
+}
+
+INSTANTIATE_TEST_SUITE_P(
+    IndexTypesAndAccumulation,
+    OpIndexPutInplaceNegativeIndexTest,
+    Combine(Values(ScalarType::Long, ScalarType::Int), Bool()));
+
+TEST_F(OpIndexPutInplaceTest, AccumulateNegativeAndPositiveAliases) {
+  TensorFactory<ScalarType::Float> tf;
+  TensorFactory<ScalarType::Long> tfl;
+  TensorFactory<ScalarType::Int> tfi;
+  Tensor values = tf.make({2, 3, 2}, {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12});
+  Tensor expected =
+      tf.make({2, 3, 2}, {6, 7, 1, 1, 5, 7, 12, 13, 1, 1, 17, 19});
+
+  for (const Tensor& index :
+       {tfl.make({3}, {-1, 2, -3}), tfi.make({3}, {-1, 2, -3})}) {
+    optional<Tensor> indices[] = {std::nullopt, index};
+    Tensor x = tf.ones({2, 3, 2});
+    op_index_put_(x, indices, values, /*accumulate=*/true);
+    EXPECT_TENSOR_EQ(x, expected);
+  }
+}

runtime/core/exec_aten/util/tensor_shape_to_c_string.cpp

@@ -48,8 +48,11 @@ std::array<char, kTensorShapeStringSizeLimit> tensor_shape_to_c_string_impl(
           static_cast<uint32_t>(elem));
     }
   }
-  *(p - 2) = ')';
-  *(p - 1) = '\0';
+  if (!shape.empty()) {
+    p -= 2;
+  }
+  *p++ = ')';
+  *p = '\0';
   return out;
 }
 } // namespace

runtime/core/exec_aten/util/test/tensor_shape_to_c_string_test.cpp

@@ -17,6 +17,14 @@ using executorch::runtime::Span;
 using executorch::runtime::tensor_shape_to_c_string;
 using executorch::runtime::internal::kMaximumPrintableTensorShapeElement;
 
+TEST(TensorShapeToCStringTest, EmptyShape) {
+  auto str32 = tensor_shape_to_c_string(Span<const std::int32_t>());
+  EXPECT_STREQ(str32.data(), "()");
+
+  auto str64 = tensor_shape_to_c_string(Span<const std::int64_t>());
+  EXPECT_STREQ(str64.data(), "()");
+}
+
 TEST(TensorShapeToCStringTest, Basic) {
   std::array<executorch::aten::SizesType, 3> sizes = {123, 456, 789};
   auto str = tensor_shape_to_c_string(