Line data Source code
1 : #if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
2 : #pragma once
3 :
4 : #include <ATen/ExpandUtils.h>
5 : #include <ATen/ScalarOps.h>
6 : #include <ATen/core/Tensor.h>
7 : #include <ATen/core/TensorBody.h>
8 : #include <c10/core/SymInt.h>
9 : #include <c10/util/irange.h>
10 : #include <optional>
11 :
12 : #ifndef AT_PER_OPERATOR_HEADERS
13 : #include <ATen/Functions.h>
14 : #include <ATen/NativeFunctions.h>
15 : #else
16 : #include <ATen/ops/alias.h>
17 : #include <ATen/ops/empty.h>
18 : #include <ATen/ops/scalar_tensor.h>
19 : #include <ATen/ops/zeros.h>
20 : #endif
21 :
22 : #include <ATen/core/List.h>
23 :
24 : #include <utility>
25 :
26 : namespace at::indexing {
27 :
28 : constexpr int64_t INDEX_MIN = c10::SymInt::min_representable_int();
29 : constexpr int64_t INDEX_MAX = -(INDEX_MIN + 1);
30 :
31 : enum class TensorIndexType { None, Ellipsis, SymInt, Boolean, Slice, Tensor };
32 :
33 : constexpr std::nullopt_t None = std::nullopt;
34 :
35 : struct TORCH_API EllipsisIndexType final {
36 : EllipsisIndexType() = default;
37 : };
38 : TORCH_API extern const EllipsisIndexType Ellipsis;
39 :
40 38 : struct TORCH_API Slice final {
41 : public:
42 116 : Slice(
43 : std::optional<c10::SymInt> start_index = std::nullopt,
44 : std::optional<c10::SymInt> stop_index = std::nullopt,
45 116 : std::optional<c10::SymInt> step_index = std::nullopt) {
46 116 : if (!step_index.has_value()) {
47 232 : step_ = c10::SymInt(1);
48 : } else {
49 0 : step_ = std::move(step_index).value();
50 : }
51 :
52 232 : TORCH_CHECK_VALUE(
53 : step_.sym_ne(0).expect_true(__FILE__, __LINE__),
54 : "slice step cannot be zero");
55 :
56 116 : if (!start_index.has_value()) {
57 348 : start_ = c10::SymInt(step_ < 0 ? INDEX_MAX : 0);
58 : } else {
59 0 : start_ = std::move(start_index).value();
60 : }
61 :
62 116 : if (!stop_index.has_value()) {
63 348 : stop_ = c10::SymInt(step_ < 0 ? INDEX_MIN : INDEX_MAX);
64 : } else {
65 0 : stop_ = std::move(stop_index).value();
66 : }
67 116 : }
68 :
69 : inline c10::SymInt start() const {
70 : return start_;
71 : }
72 :
73 : inline c10::SymInt stop() const {
74 : return stop_;
75 : }
76 :
77 : inline c10::SymInt step() const {
78 : return step_;
79 : }
80 :
81 : private:
82 : c10::SymInt start_;
83 : c10::SymInt stop_;
84 : c10::SymInt step_;
85 : };
86 :
87 : TORCH_API std::ostream& operator<<(std::ostream& stream, const Slice& slice);
88 :
89 : // `at::indexing::TensorIndex` is used for converting C++ tensor indices such as
90 : // `{None, "...", Ellipsis, 0, true, Slice(1, None, 2), torch::tensor({1, 2})}`
91 : // into its equivalent `std::vector<TensorIndex>`, so that further tensor
92 : // indexing operations can be performed using the supplied indices.
93 : //
94 : // There is one-to-one correspondence between Python and C++ tensor index types:
95 : // Python | C++
96 : // -----------------------------------------------------
97 : // `None` | `at::indexing::None`
98 : // `Ellipsis` | `at::indexing::Ellipsis`
99 : // `...` | `"..."`
100 : // `123` | `123`
101 : // `True` / `False` | `true` / `false`
102 : // `:` | `Slice()` / `Slice(None, None)`
103 : // `::` | `Slice()` / `Slice(None, None, None)`
104 : // `1:` | `Slice(1, None)`
105 : // `1::` | `Slice(1, None, None)`
106 : // `:3` | `Slice(None, 3)`
107 : // `:3:` | `Slice(None, 3, None)`
108 : // `::2` | `Slice(None, None, 2)`
109 : // `1:3` | `Slice(1, 3)`
110 : // `1::2` | `Slice(1, None, 2)`
111 : // `:3:2` | `Slice(None, 3, 2)`
112 : // `1:3:2` | `Slice(1, 3, 2)`
113 : // `torch.tensor([1, 2])`) | `torch::tensor({1, 2})`
114 : struct TORCH_API TensorIndex final {
115 : // Case 1: `at::indexing::None`
116 : TensorIndex(std::nullopt_t /*unused*/) : type_(TensorIndexType::None) {}
117 :
118 : // Case 2: "..." / `at::indexing::Ellipsis`
119 : TensorIndex(at::indexing::EllipsisIndexType /*unused*/)
120 : : type_(TensorIndexType::Ellipsis) {}
121 : TensorIndex(const char* str) : TensorIndex(at::indexing::Ellipsis) {
122 : TORCH_CHECK_VALUE(
123 : strcmp(str, "...") == 0,
124 : "Expected \"...\" to represent an ellipsis index, but got \"",
125 : str,
126 : "\"");
127 : }
128 :
129 : // Case 3: (Sym) Integer value
130 38 : TensorIndex(SymInt integer)
131 76 : : integer_(std::move(integer)), type_(TensorIndexType::SymInt) {}
132 22 : TensorIndex(int64_t integer) : TensorIndex(SymInt(integer)) {}
133 54 : TensorIndex(int integer) : TensorIndex(SymInt(integer)) {}
134 :
135 : // Case 4: Boolean value
136 : template <class T, class = std::enable_if_t<std::is_same_v<bool, T>>>
137 : TensorIndex(T boolean) : boolean_(boolean), type_(TensorIndexType::Boolean) {}
138 :
139 : // Case 5: Slice represented in `at::indexing::Slice` form
140 38 : TensorIndex(Slice slice)
141 38 : : slice_(std::move(slice)), type_(TensorIndexType::Slice) {}
142 :
143 : // Case 6: Tensor value
144 : TensorIndex(Tensor tensor)
145 : : tensor_(std::move(tensor)), type_(TensorIndexType::Tensor) {}
146 :
147 : inline bool is_none() const {
148 : return type_ == TensorIndexType::None;
149 : }
150 :
151 : inline bool is_ellipsis() const {
152 : return type_ == TensorIndexType::Ellipsis;
153 : }
154 :
155 : inline bool is_integer() const {
156 : return type_ == TensorIndexType::SymInt;
157 : }
158 :
159 : inline SymInt integer() const {
160 : return integer_;
161 : }
162 :
163 : inline bool is_boolean() const {
164 : return type_ == TensorIndexType::Boolean;
165 : }
166 :
167 : inline bool boolean() const {
168 : return boolean_;
169 : }
170 :
171 : inline bool is_slice() const {
172 : return type_ == TensorIndexType::Slice;
173 : }
174 :
175 : inline const Slice& slice() const {
176 : return slice_;
177 : }
178 :
179 : inline bool is_tensor() const {
180 : return type_ == TensorIndexType::Tensor;
181 : }
182 :
183 : inline const Tensor& tensor() const {
184 : return tensor_;
185 : }
186 :
187 : private:
188 : SymInt integer_ = 0;
189 : bool boolean_ = false;
190 : Slice slice_;
191 : Tensor tensor_;
192 : TensorIndexType type_;
193 : };
194 :
195 : TORCH_API std::ostream& operator<<(
196 : std::ostream& stream,
197 : const TensorIndex& tensor_index);
198 : TORCH_API std::ostream& operator<<(
199 : std::ostream& stream,
200 : const std::vector<TensorIndex>& tensor_indices);
201 :
202 : namespace impl {
203 : inline Tensor applySlice(
204 : const Tensor& self,
205 : int64_t dim,
206 : const c10::SymInt& start,
207 : const c10::SymInt& stop,
208 : const c10::SymInt& step,
209 : bool disable_slice_optimization,
210 : const at::Device& self_device,
211 : const std::optional<SymIntArrayRef>& self_sizes) {
212 : // TODO: implement negative step
213 : TORCH_CHECK_VALUE(
214 : step.sym_gt(0).expect_true(__FILE__, __LINE__),
215 : "step must be greater than zero");
216 :
217 : // See NOTE [nested tensor size for indexing]
218 : if (self_sizes.has_value() && !self_sizes.value().empty()) {
219 : // Skip this optimization if we are tracing, as the trace may be polymorphic
220 : // over the shape of the `self` tensor, and we still want to record
221 : // the slice.
222 : SymInt length = (self_device == at::kCPU || self_device == at::kCUDA)
223 : ? (*self_sizes)[dim]
224 : : self.sym_size(dim);
225 : if (!disable_slice_optimization &&
226 : TORCH_STATICALLY_KNOWN_TRUE(start.sym_eq(0)) &&
227 : TORCH_STATICALLY_KNOWN_TRUE(length.sym_le(stop)) && step == 1) {
228 : return self;
229 : }
230 : }
231 : return self.slice_symint(dim, start, stop, step);
232 : }
233 :
234 : inline Tensor applySelect(
235 : const Tensor& self,
236 : int64_t dim,
237 : SymInt index,
238 : int64_t real_dim,
239 : const at::Device& /*self_device*/,
240 : const std::optional<SymIntArrayRef>& self_sizes) {
241 : // See NOTE [nested tensor size for indexing]
242 : if (self_sizes.has_value()) {
243 : auto maybe_index = index.maybe_as_int();
244 : if (maybe_index.has_value()) {
245 : TORCH_CHECK_INDEX(
246 : !(maybe_index.value() == 0 && dim == 0 && self_sizes->empty()),
247 : "invalid index of a 0-dim tensor. ",
248 : "Use `tensor.item()` in Python or `tensor.item<T>()` in C++ to convert a 0-dim tensor to a number");
249 : }
250 :
251 : auto size = (*self_sizes)[dim];
252 : // Note: `size >= -index` is not equivalent to `size > -1 - index` if index
253 : // is INT64_MIN For std::numeric_limits<int64_t>::min() result of unary
254 : // minus is undefined by the standard but in practice is equal to self. On
255 : // the other hand, indexing wrapping is valid for all negative int64_t
256 : // values, as x[INT64_MIN] is the same as x[INT64_MAX]
257 : TORCH_CHECK_INDEX(
258 : size.sym_gt(-1 - index)
259 : .sym_and(size.sym_gt(index))
260 : .expect_true(__FILE__, __LINE__),
261 : "index ",
262 : index,
263 : " is out of bounds for dimension ",
264 : real_dim,
265 : " with size ",
266 : size);
267 : }
268 :
269 : // if the index is negative, do not normalize it because that would fix the
270 : // index on the current tensor size in the tracer. aten::select also works on
271 : // negative indices
272 : return self.select_symint(dim, std::move(index));
273 : }
274 :
275 : inline Tensor boolToIndexingTensorCPUOrCUDA(const Tensor& self, bool value) {
276 : // booleans add a dimension of size 1. true indexes this dimension as if 0:,
277 : // false as empty.
278 : if (value) {
279 : return at::empty({1}, self.options().dtype(kLong)).fill_(0.);
280 : } else {
281 : return at::empty({0}, self.options().dtype(kLong));
282 : }
283 : }
284 :
285 : inline Tensor boolToIndexingTensorNonNativeDeviceType(
286 : const Tensor& self,
287 : bool value) {
288 : // booleans add a dimension of size 1. true indexes this dimension as if 0:,
289 : // false as empty.
290 : if (value) {
291 : return at::zeros({1}, self.options().dtype(kLong));
292 : } else {
293 : return at::empty({0}, self.options().dtype(kLong));
294 : }
295 : }
296 :
297 : inline Tensor boolToIndexingTensor(
298 : const Tensor& self,
299 : bool value,
300 : const at::Device& self_device) {
301 : if (self_device == at::kCPU || self_device == at::kCUDA) {
302 : return boolToIndexingTensorCPUOrCUDA(self, value);
303 : } else {
304 : return boolToIndexingTensorNonNativeDeviceType(self, value);
305 : }
306 : }
307 :
308 : inline Tensor scalarToTensorNonNativeDeviceType(
309 : const Scalar& v,
310 : const TensorOptions& options) {
311 : return at::scalar_tensor(v, options);
312 : }
313 :
314 : inline void recordTensorIndex(
315 : const Tensor& tensor,
316 : std::vector<Tensor>& outIndices,
317 : int64_t* dim_ptr) {
318 : if (outIndices.empty()) {
319 : outIndices.resize(*dim_ptr + 1);
320 : outIndices[*dim_ptr] = tensor;
321 : } else {
322 : outIndices.push_back(tensor);
323 : }
324 : if (tensor.scalar_type() == kByte || tensor.scalar_type() == kBool) {
325 : *dim_ptr += tensor.dim();
326 : } else {
327 : *dim_ptr += 1;
328 : }
329 : }
330 :
331 : inline c10::List<::std::optional<Tensor>> typeConvertIndices(
332 : const Tensor& /*self*/,
333 : std::vector<Tensor> indices) {
334 : c10::List<::std::optional<Tensor>> converted_inds;
335 : converted_inds.reserve(indices.size());
336 : for (auto&& i : std::move(indices)) {
337 : converted_inds.push_back(std::move(i));
338 : }
339 : return converted_inds;
340 : }
341 :
342 : // NOTE: Why do we mirror instead of replace the `count_specified_dimensions`
343 : // function in torch/csrc/autograd/python_variable_indexing.cpp? It's because
344 : // `count_specified_dimensions` is on the hot path of Python tensor multi-dim
345 : // indexing (i.e. it's called by `applySlicing` which is called by
346 : // `THPVariable_getitem` / `THPVariable_setitem` when handling indexing of more
347 : // than one dimension). If we were to merge the Python/C++
348 : // `count_specified_dimensions` function, on the Python side we would have to
349 : // construct a `std::vector` container to be consumed by the C++
350 : // `count_specified_dimensions` function, which adds 100s of nanoseconds
351 : // overhead and is undesirable.
352 : inline int64_t count_specified_dimensions(
353 : const ArrayRef<TensorIndex>& indices) {
354 : // Count the number of indexed dimensions (everything but ellipsis and None)
355 : int64_t count = 0;
356 : for (auto& obj : indices) {
357 : if (obj.is_tensor()) {
358 : auto& tensor = obj.tensor();
359 : if (tensor.scalar_type() == kByte || tensor.scalar_type() == kBool) {
360 : count += tensor.dim();
361 : } else {
362 : count++;
363 : }
364 : } else if (!obj.is_none() && !obj.is_ellipsis() && !obj.is_boolean()) {
365 : count++;
366 : }
367 : }
368 : return count;
369 : }
370 : } // namespace impl
371 :
372 : // NOTE: Many functions below are only for consumption from Python indexing
373 : // implementation, they include:
374 : //
375 : // - `Tensor scalarToTensor(...)`
376 : // - `IntArrayRef slicePrefix1sSize(...)`
377 : // - `void copy_to(...)`
378 : // - `Tensor handleDimInMultiDimIndexing(...)`
379 : // - `Tensor dispatch_index(...)`
380 : // - `Tensor dispatch_index_put_(...)`
381 : // - `Tensor get_item(...)`
382 : // - `void set_item(...)`
383 : //
384 : // The rest of the functions are in `at::indexing::impl` namespace, signifying
385 : // that they shouldn't be used from Python indexing implementation.
386 : inline Tensor scalarToTensor(
387 : const Scalar& v,
388 : const TensorOptions& options,
389 : const at::Device& self_device) {
390 : if (self_device == at::kCPU && !v.isSymbolic()) {
391 : return at::detail::scalar_tensor_static(
392 : v,
393 : // NOLINTNEXTLINE(bugprone-unchecked-optional-access)
394 : options.dtype_opt()->toScalarType(),
395 : self_device);
396 : } else {
397 : return impl::scalarToTensorNonNativeDeviceType(v, options);
398 : }
399 : }
400 :
401 : // To match numpy semantics:
402 : // As a special case for backwards compatibility,
403 : // strip away unit dimensions from the left of 'src'
404 : inline SymIntArrayRef slicePrefix1sSize(const SymIntArrayRef& sizes) {
405 : size_t first_non1_src = sizes.size();
406 : for (const auto i : c10::irange(sizes.size())) {
407 : // Unbacked SymInt has different behavior, but this is sound because
408 : // failing to slice will only ever cause an error, not divergent
409 : // behavior
410 : if (!sizes[i].has_hint() || sizes[i] != 1) {
411 : first_non1_src = i;
412 : break;
413 : }
414 : }
415 :
416 : return sizes.slice(first_non1_src);
417 : }
418 :
419 : inline void copy_to(const Tensor& dst, const Tensor& src) {
420 : // Use TORCH_GUARD_OR_FALSE with sym_eq to avoid data-dependent-expression
421 : // errors with unbacked symbolic sizes. When we can't prove equality,
422 : // we fall through to the broadcast/expand path.
423 : auto dst_sizes = dst.sym_sizes();
424 : auto src_sizes = src.sym_sizes();
425 : bool same_sizes = dst_sizes.size() == src_sizes.size();
426 : if (same_sizes) {
427 : for (size_t i = 0; i < dst_sizes.size(); ++i) {
428 : if (!TORCH_GUARD_OR_FALSE(sym_eq(dst_sizes[i], src_sizes[i]))) {
429 : same_sizes = false;
430 : break;
431 : }
432 : }
433 : }
434 : if (same_sizes) {
435 : // A shortcut to avoid generating hard-coded constant sizes during tracing.
436 : // This is not a perfect solution: when src & dst have different shapes,
437 : // constants will still appear. Users can workaround that case by
438 : // dst[index..] = src.reshape(..)
439 : dst.copy_(src);
440 : return;
441 : } else if (src.dim() == 0 && src.device().type() == at::kCPU) {
442 : dst.fill_(src);
443 : return;
444 : }
445 : auto src_view = src.view_symint(slicePrefix1sSize(src.sym_sizes()));
446 : c10::MaybeOwned<Tensor> b_src = expand_inplace(dst, src_view, "setitem");
447 : dst.copy_(*b_src);
448 : }
449 :
450 : // See NOTE [ Setting `disable_slice_optimization` when calling C++ tensor
451 : // indexing functions from Python ]
452 : inline Tensor handleDimInMultiDimIndexing(
453 : const Tensor& prev_dim_result,
454 : const Tensor& original_tensor,
455 : const TensorIndex& index,
456 : int64_t* dim_ptr,
457 : int64_t* specified_dims_ptr,
458 : int64_t real_dim,
459 : std::vector<Tensor>& outIndices,
460 : bool disable_slice_optimization,
461 : const at::Device& original_tensor_device,
462 : const std::optional<SymIntArrayRef>& prev_dim_result_sizes) {
463 : if (index.is_integer()) {
464 : return impl::applySelect(
465 : prev_dim_result,
466 : *dim_ptr,
467 : index.integer(),
468 : real_dim,
469 : original_tensor_device,
470 : prev_dim_result_sizes);
471 : } else if (index.is_slice()) {
472 : Tensor result = impl::applySlice(
473 : prev_dim_result,
474 : *dim_ptr,
475 : index.slice().start(),
476 : index.slice().stop(),
477 : index.slice().step(),
478 : /*disable_slice_optimization=*/disable_slice_optimization,
479 : original_tensor_device,
480 : prev_dim_result_sizes);
481 : (*dim_ptr)++;
482 : if (!outIndices.empty()) {
483 : outIndices.resize(outIndices.size() + 1);
484 : }
485 : return result;
486 : } else if (index.is_ellipsis()) {
487 : auto ellipsis_ndims = original_tensor.dim() - *specified_dims_ptr;
488 : (*dim_ptr) += ellipsis_ndims;
489 : if (!outIndices.empty()) {
490 : outIndices.resize(outIndices.size() + ellipsis_ndims);
491 : }
492 : return prev_dim_result;
493 : } else if (index.is_none()) {
494 : Tensor result = prev_dim_result.unsqueeze(*dim_ptr);
495 : (*dim_ptr)++;
496 : if (!outIndices.empty()) {
497 : outIndices.resize(outIndices.size() + 1);
498 : }
499 : return result;
500 : } else if (index.is_boolean()) {
501 : Tensor result = prev_dim_result.unsqueeze(*dim_ptr);
502 : impl::recordTensorIndex(
503 : impl::boolToIndexingTensor(
504 : result, index.boolean(), original_tensor_device),
505 : outIndices,
506 : dim_ptr);
507 : return result;
508 : } else if (index.is_tensor()) {
509 : Tensor result = prev_dim_result;
510 : const Tensor& tensor = index.tensor();
511 : auto scalar_type = tensor.scalar_type();
512 : bool is_batched = tensor.key_set().has_any(c10::DispatchKeySet(
513 : {c10::DispatchKey::FuncTorchBatched,
514 : c10::DispatchKey::BatchedNestedTensor}));
515 : if (tensor.dim() == 0 && !is_batched &&
516 : at::isIntegralType(scalar_type, /*includeBool=*/true)) {
517 : if (scalar_type != at::kByte && scalar_type != at::kBool) {
518 : result = impl::applySelect(
519 : result,
520 : *dim_ptr,
521 : tensor.item().toSymInt(),
522 : real_dim,
523 : original_tensor_device,
524 : prev_dim_result_sizes);
525 : } else {
526 : result = result.unsqueeze(*dim_ptr);
527 : if (scalar_type == at::kBool) {
528 : impl::recordTensorIndex(
529 : impl::boolToIndexingTensor(
530 : result, tensor.item<bool>() != 0, original_tensor_device),
531 : outIndices,
532 : dim_ptr);
533 : } else {
534 : impl::recordTensorIndex(
535 : impl::boolToIndexingTensor(
536 : result, tensor.item<uint8_t>() != 0, original_tensor_device),
537 : outIndices,
538 : dim_ptr);
539 : }
540 : }
541 : } else {
542 : impl::recordTensorIndex(tensor, outIndices, dim_ptr);
543 : }
544 : return result;
545 : } else {
546 : TORCH_INTERNAL_ASSERT(false, "Invalid TensorIndex type");
547 : }
548 : }
549 :
550 : namespace impl {
551 : // This mirrors `applySlicing` in
552 : // torch/csrc/autograd/python_variable_indexing.cpp
553 : inline Tensor applySlicing(
554 : const Tensor& self,
555 : const ArrayRef<TensorIndex>& indices,
556 : std::vector<Tensor>& outIndices,
557 : bool disable_slice_optimization,
558 : const at::Device& self_device,
559 : const std::optional<SymIntArrayRef>& self_sizes) {
560 : int64_t dim = 0;
561 : int64_t specified_dims = impl::count_specified_dimensions(indices);
562 :
563 : // See NOTE [nested tensor size for indexing]
564 : if (self_sizes.has_value()) {
565 : TORCH_CHECK_INDEX(
566 : specified_dims <= (int64_t)self_sizes->size(),
567 : "too many indices for tensor of dimension ",
568 : (int)self_sizes->size());
569 : }
570 :
571 : Tensor result = self;
572 : for (const auto i : c10::irange(indices.size())) {
573 : auto& obj = indices[i];
574 : // See NOTE [nested tensor size for indexing]
575 : std::optional<SymIntArrayRef> result_sizes = result.is_nested()
576 : ? std::optional<SymIntArrayRef>(std::nullopt)
577 : : std::optional<SymIntArrayRef>(result.sym_sizes());
578 : result = handleDimInMultiDimIndexing(
579 : /*prev_dim_result=*/result,
580 : /*original_tensor=*/self,
581 : /*index=*/obj,
582 : /*dim_ptr=*/&dim,
583 : /*specified_dims_ptr=*/&specified_dims,
584 : /*real_dim=*/static_cast<int64_t>(i),
585 : /*outIndices=*/outIndices,
586 : /*disable_slice_optimization=*/disable_slice_optimization,
587 : /*original_tensor_device=*/self_device,
588 : /*prev_dim_result_sizes=*/result_sizes);
589 : }
590 : return result;
591 : }
592 : } // namespace impl
593 :
594 : inline Tensor dispatch_index(const Tensor& self, std::vector<Tensor> indices) {
595 : // Remove trailing null elements from indices
596 : while (!indices.empty() && !indices.back().defined()) {
597 : indices.pop_back();
598 : }
599 : return self.index(impl::typeConvertIndices(self, std::move(indices)));
600 : }
601 :
602 : inline Tensor dispatch_index_put_(
603 : Tensor& self,
604 : std::vector<Tensor> indices,
605 : const Tensor& value) {
606 : // Remove trailing null elements from indices
607 : while (!indices.empty() && !indices.back().defined()) {
608 : indices.pop_back();
609 : }
610 : return self.index_put_(
611 : impl::typeConvertIndices(self, std::move(indices)), value);
612 : }
613 :
614 : // NOTE [ Setting `disable_slice_optimization` when calling C++ tensor indexing
615 : // functions from Python ]
616 : //
617 : // Question: When should we set `disable_slice_optimization` to `true` when
618 : // calling C++ tensor indexing functions from Python indexing code?
619 : //
620 : // Answer: What "slice optimization" means: when we have a slicing expression
621 : // like `x[0:5, 0]`, where the sliced tensor was of size 5 in dimension 0, we
622 : // would skip dispatching the actual slice call as an optimization. However,
623 : // here are the cases where we DON'T want this optimization:
624 : //
625 : // 1. When we are doing 1-D slicing (e.g. `tensor[:]`).
626 : // Reason: we always return a shallow copy for expressions such as
627 : // `tensor[:]` / `tensor[...]` / `tensor[:, :]`. (Note that for `tensor[:,
628 : // :]`, we return an alias of `tensor` by doing the following:
629 : // ```
630 : // Tensor sliced = impl::applySlicing(self, indices, tensorIndices,
631 : // disable_slice_optimization, self_device, self_sizes); if
632 : // (tensorIndices.empty()) {
633 : // if (sliced.is_same(self)) {
634 : // // ensure we return a shallow copy for things like x[...]
635 : // sliced = at::alias(sliced);
636 : // }
637 : // return sliced;
638 : // }
639 : // ```)
640 : // 2. When we are doing JIT tracing.
641 : // Reason: JIT tracing needs the `self.slice(...)` call to properly trace the
642 : // slice operation.
643 :
644 : // This mirrors `THPVariable_getitem` in
645 : // torch/csrc/autograd/python_variable_indexing.cpp See NOTE [ Setting
646 : // `disable_slice_optimization` when calling C++ tensor indexing functions from
647 : // Python ]
648 : inline Tensor get_item(
649 : const Tensor& self,
650 : const ArrayRef<TensorIndex>& indices,
651 : bool disable_slice_optimization = false) {
652 : at::Device self_device = self.device();
653 : // NOTE [nested tensor size for indexing]
654 : // nested tensor does not have a size (yet) so for now we represent its size
655 : // as null may need to be changed after we reach a better solution for nested
656 : // tensor size
657 : std::optional<SymIntArrayRef> self_sizes = self.is_nested()
658 : ? std::optional<SymIntArrayRef>(std::nullopt)
659 : : std::optional<SymIntArrayRef>(self.sym_sizes());
660 :
661 : // handle simple types: integers, slices, none, ellipsis, bool
662 : if (indices.size() == 1) {
663 : const TensorIndex& index = indices[0];
664 : if (index.is_integer()) {
665 : return impl::applySelect(
666 : self, 0, index.integer(), 0, self_device, self_sizes);
667 : } else if (index.is_slice()) {
668 : return impl::applySlice(
669 : self,
670 : 0,
671 : index.slice().start(),
672 : index.slice().stop(),
673 : index.slice().step(),
674 : /*disable_slice_optimization=*/true,
675 : self_device,
676 : self_sizes);
677 : } else if (index.is_none()) {
678 : return self.unsqueeze(0);
679 : } else if (index.is_ellipsis()) {
680 : return at::alias(self);
681 : } else if (index.is_boolean()) {
682 : Tensor result = self.unsqueeze(0);
683 : return dispatch_index(
684 : result,
685 : std::vector<Tensor>{impl::boolToIndexingTensor(
686 : result, index.boolean(), self_device)});
687 : }
688 : }
689 :
690 : std::vector<Tensor> tensorIndices;
691 : Tensor sliced = impl::applySlicing(
692 : self,
693 : indices,
694 : tensorIndices,
695 : disable_slice_optimization,
696 : self_device,
697 : self_sizes);
698 : if (tensorIndices.empty()) {
699 : if (sliced.is_same(self)) {
700 : // ensure we return a shallow copy for things like x[...]
701 : sliced = at::alias(sliced);
702 : }
703 : return sliced;
704 : }
705 :
706 : // indexing by tensors ("advanced" indexing)
707 : return dispatch_index(sliced, std::move(tensorIndices));
708 : }
709 :
710 : // This mirrors `THPVariable_setitem` in
711 : // torch/csrc/autograd/python_variable_indexing.cpp for "the assigned value is a
712 : // Tensor" case See NOTE [ Setting `disable_slice_optimization` when calling C++
713 : // tensor indexing functions from Python ]
714 : inline void set_item(
715 : const Tensor& self,
716 : const ArrayRef<TensorIndex>& indices,
717 : const Tensor& value,
718 : bool disable_slice_optimization = false) {
719 : at::Device self_device = self.device();
720 : SymIntArrayRef self_sizes = self.sym_sizes();
721 :
722 : // handle simple types: integers, slices, ellipsis, bool
723 : if (indices.size() == 1) {
724 : const TensorIndex& index = indices[0];
725 : if (index.is_boolean() && !index.boolean()) {
726 : // do nothing for false (technically we should check the size, but we
727 : // don't have real 0-sized shapes.
728 : return;
729 : } else if (index.is_ellipsis()) {
730 : copy_to(self, value);
731 : return;
732 : } else if (index.is_none() || (index.is_boolean() && index.boolean())) {
733 : copy_to(self.unsqueeze(0), value);
734 : return;
735 : } else if (index.is_integer()) {
736 : copy_to(
737 : impl::applySelect(
738 : self, 0, index.integer(), 0, self_device, self_sizes),
739 : value);
740 : return;
741 : } else if (index.is_slice()) {
742 : copy_to(
743 : impl::applySlice(
744 : self,
745 : 0,
746 : index.slice().start(),
747 : index.slice().stop(),
748 : index.slice().step(),
749 : /*disable_slice_optimization=*/disable_slice_optimization,
750 : self_device,
751 : self_sizes),
752 : value);
753 : return;
754 : }
755 : }
756 :
757 : std::vector<Tensor> tensorIndices;
758 : Tensor sliced = impl::applySlicing(
759 : self,
760 : indices,
761 : tensorIndices,
762 : disable_slice_optimization,
763 : self_device,
764 : self_sizes);
765 : if (tensorIndices.empty()) {
766 : copy_to(sliced, value);
767 : return;
768 : }
769 :
770 : SymIntArrayRef valueSizes = value.sym_sizes();
771 : SymIntArrayRef slicedValueSizes = slicePrefix1sSize(valueSizes);
772 : Tensor valuesSliced;
773 : if (!valueSizes.equals(slicedValueSizes)) {
774 : valuesSliced = value.view_symint(slicedValueSizes);
775 : } else {
776 : valuesSliced = value;
777 : }
778 : dispatch_index_put_(sliced, std::move(tensorIndices), valuesSliced);
779 : return;
780 : }
781 :
782 : } // namespace at::indexing
783 :
784 : #else
785 : #error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
786 : #endif // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
|
