Function ConvertToIndexingTensor

oneflow/api/python/functional/indexing.cpp:158–194

}

Maybe<Tensor> ConvertToIndexingTensor(PyObject* object) {
  // NOTE: convert data to indexing will ensure in eager mode
  LazyMode::Guard lazy_mode_disabled_guard(/*is_enabled*/ false);
  const DataType dtype = InferScalarType(object);
  const auto& device = JUST(Device::New("cpu"));

  // index type must be integers
  if (!(IsIntegralDataType(dtype) || (IsBoolDataType(dtype)))) {
    return Error::IndexError() << "only integers, slices (`:`), ellipsis (`...`), numpy.newaxis "
                                  "(`None`) and integer or boolean arrays are valid indices";
  }
  // In advanced indexing condition, index can be array object, need to handle it specially.
  if (PyArray_Check(object)) {
    return TensorWithData(object, NullOpt, device, /*requires_grad=*/false, /*pin_memory=*/false);
  }

  const auto& sizes = InferArraySizes(object);
  const auto& tensor = JUST(functional::Empty(sizes, CHECK_JUST(DType::Get(dtype)), device,
                                              /*requires_grad=*/false, /*pin_memory=*/false));
  // Prevent the python object release until the callback is complete.
  Py_INCREF(object);
  auto handle = std::shared_ptr<PyObject>(PyObjectPtr(object));

  JUST(PhysicalRun([&](InstructionsBuilder* builder) -> Maybe<void> {
    return builder->AccessBlobByCallback(
        JUST(tensor->AsLocalTensor()),
        [handle](ep::Stream* stream,
                 const std::shared_ptr<vm::EagerBlobObject>& eager_blob_object) {
          CHECK_JUST(Singleton<ForeignLockHelper>::Get()->WithScopedAcquire([&]() -> Maybe<void> {
            ParseArrayToTensor(handle.get(), eager_blob_object);
            return Maybe<void>::Ok();
          }));
        },
        "mut");
  }));
  return tensor;
}

IndexItem UnpackIndexItem(PyObject* object) {
  if (object == Py_Ellipsis) {