python_variable.cpp

#include <ATen/DTensorState.h>
#include <ATen/NamedTensorUtils.h>
#include <ATen/native/Resize.h>
#include <c10/core/DeviceType.h>
#include <c10/core/SymIntArrayRef.h>
#include <c10/core/impl/GPUTrace.h>
#include <c10/core/impl/HermeticPyObjectTLS.h>
#include <c10/core/impl/PythonDispatcherTLS.h>
#include <c10/util/FbcodeMaps.h>
#include <c10/util/SmallVector.h>
#include <c10/util/irange.h>
#include <pybind11/pytypes.h>
#include <torch/csrc/Device.h>
#include <torch/csrc/DynamicTypes.h>
#include <torch/csrc/Exceptions.h>
#include <torch/csrc/PyInterpreter.h>
#include <torch/csrc/Size.h>
#include <torch/csrc/THP.h>
#include <torch/csrc/Types.h>
#include <torch/csrc/autograd/autograd.h>
#include <torch/csrc/autograd/edge.h>
#include <torch/csrc/autograd/function.h>
#include <torch/csrc/autograd/python_cpp_function.h>
#include <torch/csrc/autograd/python_hook.h>
#include <torch/csrc/autograd/python_torch_functions.h>
#include <torch/csrc/autograd/python_variable_indexing.h>
#include <torch/csrc/autograd/utils/error_messages.h>
#include <torch/csrc/autograd/utils/wrap_outputs.h>
#include <torch/csrc/autograd/variable.h>
#include <torch/csrc/distributed/Placement.h>
#include <torch/csrc/jit/frontend/tracer.h>
#include <torch/csrc/jit/python/pybind_utils.h>
#include <torch/csrc/tensor/python_tensor.h>
#include <torch/csrc/utils/pybind.h>
#include <torch/csrc/utils/pycfunction_helpers.h>
#include <torch/csrc/utils/pyobject_preservation.h>
#include <torch/csrc/utils/python_arg_parser.h>
#include <torch/csrc/utils/python_compat.h>
#include <torch/csrc/utils/python_dispatch.h>
#include <torch/csrc/utils/python_strings.h>
#include <torch/csrc/utils/tensor_new.h>
#include <torch/csrc/utils/tensor_numpy.h>

#include <torch/csrc/utils/torch_dispatch_mode.h>

#include <ATen/ATen.h>

#include <structmember.h>
#include <cstdint>
#include <memory>
#include <sstream>
#include <utility>
#include <vector>

using namespace at;
using namespace torch;
using namespace torch::autograd;
using torch::utils::PyObjectPreservation;

namespace {
class OperatorArgsKwargsView {
 public:
  OperatorArgsKwargsView(
      const c10::OperatorHandle& op,
      const std::vector<c10::IValue>& arguments);
  using args_iterator = const c10::IValue*;

  args_iterator args_begin() const {
    return arguments_.data();
  }

  args_iterator args_end() const {
    return arguments_.data() + positional_default_start_;
  }

  auto num_positional_args() const {
    return positional_default_start_;
  }

  auto kwarg_start_index() const {
    return first_non_default_kwarg_;
  }

  struct kwargs_iterator {
    kwargs_iterator() = default;
    kwargs_iterator(const OperatorArgsKwargsView* parent, size_t current)
        : parent_(parent), current_(current) {}

    kwargs_iterator(const kwargs_iterator&) = default;
    kwargs_iterator& operator=(const kwargs_iterator&) = default;

    kwargs_iterator& operator++() {
      do {
        current_++;
      } while (current_ < parent_->arguments_.size() &&
               parent_->is_default(current_));
      return *this;
    }

    kwargs_iterator operator++(int) {
      auto copy = *this;
      ++(*this);
      return copy;
    }

    const c10::IValue& operator*() const {
      return parent_->arguments_[current_];
    }

    const c10::IValue* operator->() const {
      return &operator*();
    }

    int64_t underlying_index() const {
      return current_;
    }

    bool operator==(const kwargs_iterator& rhs) const {
      return parent_ == rhs.parent_ && current_ == rhs.current_;
    }

    bool operator!=(const kwargs_iterator& rhs) {
      return !(*this == rhs);
    }

   private:
    const OperatorArgsKwargsView* parent_ = nullptr;
    size_t current_ = 0;
  };

  kwargs_iterator kwargs_begin() const {
    return kwargs_iterator(this, first_non_default_kwarg_);
  }

  kwargs_iterator kwargs_end() const {
    return kwargs_iterator(this, arguments_.size());
  }

 private:
  bool is_default(size_t idx) const {
    const auto& arg = op_.schema().arguments()[idx];
    if (!arg.default_value().has_value()) {
      return false;
    }
    const auto& default_ivalue = *arg.default_value();
    const auto& ivalue = arguments_[idx];
    if (default_ivalue != ivalue) {
      return false;
    }
    return true;
  }

  const c10::OperatorHandle& op_;
  c10::ArrayRef<c10::IValue> arguments_;
  // About all the pointers:
  //
  // f(int x, int y = 0, *, int z = 0)
  //                                  ^- arguments.size()
  //                        ^- kwarg_only_start
  //          ^- positional_default_start
  //   ^- 0
  int64_t positional_default_start_;
  int64_t first_non_default_kwarg_;
};

OperatorArgsKwargsView::OperatorArgsKwargsView(
    const c10::OperatorHandle& op,
    const std::vector<c10::IValue>& arguments)
    : op_(op), arguments_(arguments) {
  // Find the split point between kwarg-only and regular.  Since most functions
  // don't have kwarg-only arguments, it is more efficient to scan from the
  // right (but ideally, this would just be precomputed in FunctionSchema
  // itself).  (NB: minus one in the loop is because we're testing if the
  // *next* argument is kwarg-only before we advance the starting index)
  const int64_t signed_arguments_size = static_cast<int64_t>(arguments.size());
  int64_t kwarg_only_start = signed_arguments_size;
  for (; kwarg_only_start > 0; kwarg_only_start--) {
    const auto& arg = op.schema().arguments()[kwarg_only_start - 1];
    if (!arg.kwarg_only()) {
      break;
    }
  }

  // Find the first positional argument that isn't defaulted
  positional_default_start_ = kwarg_only_start;
  for (; positional_default_start_ > 0; positional_default_start_--) {
    if (!is_default(positional_default_start_ - 1)) {
      break;
    }
  }

  // kwargs_iterator will skip default kwargs when incremented, but we
  // need to skip any initial run of default kwargs ourselves.
  first_non_default_kwarg_ = kwarg_only_start;
  for (; first_non_default_kwarg_ < signed_arguments_size;
       ++first_non_default_kwarg_) {
    if (!is_default(first_non_default_kwarg_)) {
      break;
    }
  }
}
} // namespace

std::pair<py::object, py::dict> parseIValuesToPyArgsKwargs(
    const c10::OperatorHandle& op,
    const std::vector<c10::IValue>& arguments) {
  TORCH_CHECK(
      PyGILState_Check(),
      "GIL must be held before you call parseIValuesToPyArgsKwargs");
  const auto& schema = op.schema();
  py::dict kwargs;

  OperatorArgsKwargsView args_kwargs(op, arguments);
  auto args = py::reinterpret_steal<py::object>(
      PyTuple_New(args_kwargs.num_positional_args()));

  auto schemaAwareToPyObject =
      [&schema](size_t idx, const c10::IValue& argument) -> py::object {
    const auto& arg = schema.arguments()[idx];
    auto match = [&](c10::TypeKind kind) {
      const auto& t = arg.real_type();
      if (t->kind() == kind)
        return true;
      if (auto opt_t = t->cast<c10::OptionalType>()) {
        if (opt_t->getElementType()->kind() == kind)
          return true;
      }
      return false;
    };
    if (argument.isNone()) {
      return py::none();
    } else if (match(c10::ScalarTypeType::Kind)) {
      auto* obj = getTHPDtype(static_cast<c10::ScalarType>(argument.toInt()));
      return py::reinterpret_borrow<py::object>(
          reinterpret_cast<PyObject*>(obj));
    } else if (match(c10::LayoutType::Kind)) {
      auto* obj = getTHPLayout(static_cast<c10::Layout>(argument.toInt()));
      return py::reinterpret_borrow<py::object>(
          reinterpret_cast<PyObject*>(obj));
    } else if (match(c10::MemoryFormatType::Kind)) {
      return py::cast(static_cast<c10::MemoryFormat>(argument.toInt()));
    } else {
      return torch::jit::toPyObject(argument);
    }
  };

  // Populate positional arguments
  size_t idx = 0;
  for (auto argument_it = args_kwargs.args_begin();
       argument_it != args_kwargs.args_end();
       ++argument_it) {
    PyTuple_SET_ITEM(
        args.ptr(),
        idx,
        schemaAwareToPyObject(idx, *argument_it).release().ptr());
    idx++;
  }

  // Populate keyword arguments
  for (auto argument_it = args_kwargs.kwargs_begin();
       argument_it != args_kwargs.kwargs_end();
       ++argument_it) {
    const auto& arg = schema.arguments()[argument_it.underlying_index()];
    kwargs[py::cast(arg.name())] =
        schemaAwareToPyObject(argument_it.underlying_index(), *argument_it);
  }
  return std::make_pair(std::move(args), std::move(kwargs));
}

void pushPyOutToStack(
    const c10::OperatorHandle& op,
    torch::jit::Stack* stack,
    py::object out,
    const char* msg) {
  TORCH_CHECK(
      PyGILState_Check(), "GIL must be held before you call pushPyOutToStack");
  const auto& schema_returns = op.schema().returns();
  const auto num_returns = schema_returns.size();
  if (num_returns == 0) {
    // Check that we got a None return from Python. Anything else is an error.
    TORCH_CHECK(
        out.is_none(),
        "Expected ",
        msg,
        " for ",
        op.operator_name(),
        " to return None but it returned something else instead.");
  } else if (num_returns == 1) {
    torch::jit::push(
        stack, torch::jit::toIValue(out.ptr(), schema_returns[0].real_type()));
  } else {
    auto outs = py::cast<py::sequence>(out);
    for (const auto idx : c10::irange(outs.size())) {
      torch::jit::push(
          stack,
          torch::jit::toIValue(
              outs[idx].ptr(), schema_returns[idx].real_type()));
    }
  }
}

namespace {

c10::TensorImpl::SizesStridesPolicy parseSizesStridesPolicyArgument(
    std::string_view arg) {
  if (arg == "strides") {
    return c10::TensorImpl::SizesStridesPolicy::CustomStrides;
  }

  if (arg == "sizes") {
    return c10::TensorImpl::SizesStridesPolicy::CustomSizes;
  }

  TORCH_CHECK_VALUE(
      false,
      "Unknown sizes_strides_policy: ",
      arg,
      "; expected 'strides' or 'sizes'");
}
} // anonymous namespace

PyObject* THPVariableClass = nullptr;

PyObject* ParameterClass = nullptr;

// clang-tidy gets confused by static const
static constexpr const char* VOLATILE_WARNING =
    "volatile was removed and now has no effect. Use "
    "`with torch.no_grad():` instead.";

static void TORCH_CHECK_TENSOR_SUBTYPE(PyObject* cls);

static bool check_has_torch_dispatch(PyObject* obj) {
  if (THPVariable_CheckExact(obj)) {
    return false;
  }
  py::object attr = PyObject_FastGetAttrString(obj, "__torch_dispatch__");
  return (
      attr.ptr() != nullptr &&
      attr.ptr() != torch::disabled_torch_dispatch_impl());
}

// NOLINTNEXTLINE(*-c-arrays,cppcoreguidelines-avoid-non-const-global-variables)
static PyObject* device_to_py_class_[static_cast<size_t>(
    c10::DeviceType::COMPILE_TIME_MAX_DEVICE_TYPES)];

void registerPythonTensorClass(
    const std::string& device,
    PyObject* python_tensor_class) {
  c10::Device dev(device);

  TORCH_CHECK(
      dev.type() == kXLA, "Only the python class for XLA can be overridden");
  if (device_to_py_class_[static_cast<size_t>(dev.type())] != nullptr) {
    TORCH_WARN(
        "Overriding a previously registered python class for ", dev.str());
  }

  device_to_py_class_[static_cast<size_t>(dev.type())] = python_tensor_class;
}

static PyObject* getPythonTensorClass(c10::Device d) {
  return device_to_py_class_[static_cast<size_t>(d.type())];
}

void activateGPUTrace() {
  c10::impl::GPUTrace::set_trace(getPyInterpreter());
}

static void check_tensor_subclass(PyObject* obj, PyTypeObject* type) {
  TORCH_CHECK(
      PyObject_TypeCheck(obj, type),
      "Creating a new Tensor subclass ",
      type->tp_name,
      " but the raw Tensor object is already associated to a python object ",
      "of type ",
      Py_TYPE(obj)->tp_name,
      " which is not a subclass of the requested type");
}

// Generic for const Tensor& or Tensor&&
template <typename T>
static PyObject* THPVariable_WrapWithType(
    T&& var,
    std::optional<PyTypeObject*> desired_type) {
  if (!var.defined()) {
    Py_RETURN_NONE;
  }

  c10::TensorImpl* tensor_impl = var.unsafeGetTensorImpl();
  c10::impl::PyObjectSlot* pyobj_slot = tensor_impl->pyobj_slot();

  PyObject* obj = pyobj_slot->load_pyobj();
  if (obj) {
    if (desired_type) {
      check_tensor_subclass(obj, *desired_type);
    }
    return Py_NewRef(obj);
  }

  PyTypeObject* type = reinterpret_cast<PyTypeObject*>(THPVariableClass);
  if (desired_type) {
    type = *desired_type;
  } else if (C10_UNLIKELY(var.device().type() == c10::kXLA)) {
    if (auto clazz = getPythonTensorClass(var.device())) {
      type = reinterpret_cast<PyTypeObject*>(clazz);
    }
  }

  obj = type->tp_alloc(type, 0);
  TORCH_CHECK(obj, "Failed to allocate a ", type->tp_name, " object");

  // Ensure that PyUnstable_TryIncref calls don't fail spuriously in
  // free-threaded Python.
  PyUnstable_EnableTryIncRef(obj);

  auto v = reinterpret_cast<THPVariable*>(obj);
  new (&v->cdata) Tensor(std::forward<T>(var));

  if (THPVariable_Unpack(obj).is_uniquely_owned()) {
    // We can use a faster non-atomic code path if we have the only reference to
    // a fresh Tensor.
    PyObjectPreservation::init_fresh_nonatomic(tensor_impl, pyobj_slot, obj);
    return obj;
  }

  PyObject* wrapper =
      PyObjectPreservation::init_once(tensor_impl, pyobj_slot, obj);
  if (wrapper != obj) {
    // Another thread beat us to it
    Py_DECREF(obj);
    if (desired_type) {
      check_tensor_subclass(wrapper, *desired_type);
    }
    return Py_NewRef(wrapper);
  }
  return obj;
}

PyObject* THPVariable_Wrap(at::TensorBase&& var) {
  return THPVariable_WrapWithType(std::move(var), std::nullopt);
}

PyObject* THPVariable_Wrap(const at::TensorBase& var) {
  return THPVariable_WrapWithType(var, std::nullopt);
}

PyObject* THPVariable_Wrap(const at::TensorBase& var, PyTypeObject* type) {
  return THPVariable_WrapWithType(var, type);
}

static PyObject* THPVariable_pynew(
    PyTypeObject* type,
    PyObject* args,
    PyObject* kwargs);

static PyObject* THPVariable_fix_weakref(PyObject* self, PyObject* noargs) {
  const auto& var = THPVariable_Unpack(self);
  Py_DECREF(THPVariable_Wrap(var));
  Py_RETURN_NONE;
}

// Maps the given python callable over a vector of items, returning a vector
// of the same type of items.
template <typename T>
static std::vector<T> map_py_func(
    const py::function& func,
    const std::vector<T>& items) {
  std::vector<T> new_items;
  new_items.reserve(items.size());
  for (auto& item : items) {
    new_items.emplace_back(py::cast<T>(func(item)));
  }
  return new_items;
}

template <>
std::vector<at::Tensor> map_py_func(
    const py::function& func,
    const std::vector<at::Tensor>& items) {
  std::vector<at::Tensor> new_items;
  new_items.reserve(items.size());
  for (auto& item : items) {
    auto output = func(item);
    if (output.is(py::none())) {
      // treat None value as an undefined tensor
      new_items.emplace_back();
    } else {
      new_items.emplace_back(py::cast<at::Tensor>(output));
    }
  }
  return new_items;
}

static PyObject* view_func_impl(
    PyObject* _self,
    PyObject* args,
    PyObject* kwargs,
    bool check_has_same_meta) {
  HANDLE_TH_ERRORS
  const auto& self = THPVariable_Unpack(_self);

  static PythonArgParser parser({
      "_view_func(Tensor new_base, PyObject* symint_visitor_fn=None, PyObject* tensor_visitor_fn=None)",
  });
  ParsedArgs<3> parsed_args{};
  auto r = parser.parse(_self, args, kwargs, parsed_args);
  auto new_base = r.tensor(0);
  PyObject* symint_visitor_fn = r.pyobject(1);
  PyObject* tensor_visitor_fn = r.pyobject(2);

  // Ensure that self is indeed a backward differentiable view
  // If not, we return an undefined Tensor (None) and let the user handle it.
  auto diff_view_meta = torch::autograd::impl::get_view_autograd_meta(self);
  at::Tensor out;
  if (diff_view_meta && diff_view_meta->has_bw_view()) {
    const auto& view_info = diff_view_meta->get_backward_view();
    // Ensure that the newly provided base is similar to the original base
    if (!check_has_same_meta ||
        torch::autograd::utils::has_same_meta(new_base, view_info.base_)) {
      // Do the actual view replay
      if (view_info.has_view_fn()) {
        auto& view_func = view_info.view_fn();

        // Determine new SymInt / tensor state as needed.
        std::optional<std::vector<c10::SymInt>> new_symints = std::nullopt;
        if (symint_visitor_fn != Py_None) {
          new_symints = map_py_func(
              py::cast<py::function>(symint_visitor_fn),
              view_func.get_symints());
        }

        std::optional<std::vector<at::Tensor>> new_tensors = std::nullopt;
        if (tensor_visitor_fn != Py_None) {
          new_tensors = map_py_func(
              py::cast<py::function>(tensor_visitor_fn),
              view_func.get_tensors());
        }

        // call view func
        if (new_symints.has_value() || new_tensors.has_value()) {
          out = (*view_func.clone_and_set(new_symints, new_tensors))(new_base);
        } else {
          out = view_func(new_base);
        }
      } else {
        out = new_base.as_strided(
            self.sizes(), self.strides(), self.storage_offset());
      }
    }
  }
  return THPVariable_Wrap(out);
  END_HANDLE_TH_ERRORS
}

static PyObject* THPVariable_view_func(
    PyObject* self_,
    PyObject* args,
    PyObject* kwargs) {
  return view_func_impl(self_, args, kwargs, /*check_has_same_meta=*/true);
}

static PyObject* THPVariable_view_func_unsafe(
    PyObject* self_,
    PyObject* args,
    PyObject* kwargs) {
  return view_func_impl(self_, args, kwargs, /*check_has_same_meta=*/false);
}

static PyObject* rev_view_func_impl(PyObject* self_, PyObject* arg) {
  HANDLE_TH_ERRORS
  const auto& self = THPVariable_Unpack(self_);
  TORCH_CHECK(
      THPVariable_Check(arg),
      "_rev_view_func expect a single argument that is a Tensor");
  const auto& new_view = THPVariable_Unpack(arg);

  // Ensure that self is indeed a backward differentiable view
  // If not, we return an undefined Tensor (None) and let the user handle it.
  auto diff_view_meta = torch::autograd::impl::get_view_autograd_meta(self);
  at::Tensor out;
  if (diff_view_meta && diff_view_meta->has_bw_view()) {
    const auto& view_info = diff_view_meta->get_backward_view();
    // Do the actual view replay
    TORCH_CHECK(view_info.has_view_fn(), "No _rev_view_func() found");
    out = view_info.rev_view_fn()(new_view);
  }
  return THPVariable_Wrap(out);
  END_HANDLE_TH_ERRORS
}

static PyObject* THPVariable_rev_view_func_unsafe(
    PyObject* self_,
    PyObject* arg) {
  return rev_view_func_impl(self_, arg);
}

// Instantiates a subclass of self with the same data.
static PyObject* THPVariable_as_subclass(
    PyObject* _self,
    PyObject* args,
    PyObject* kwargs) {
  HANDLE_TH_ERRORS
  const auto& self = THPVariable_Unpack(_self);
  static PythonArgParser parser({
      "as_subclass(PyObject* cls)",
  });
  ParsedArgs<1> parsed_args{};
  auto r = parser.parse(_self, args, kwargs, parsed_args);
  PyObject* cls = r.pyobject(0);
  TORCH_CHECK_TENSOR_SUBTYPE(cls);
  // guard completely turns off torch dispatch modes, doesn't just pop off the
  // stack
  torch_dispatch_mode::StashTorchDispatchStackGuard td_g;
  c10::impl::DisablePythonDispatcher dpd_g;
  PyObject* obj = THPVariable_WrapWithType(self.alias(), (PyTypeObject*)cls);
  if (check_has_torch_dispatch(obj)) {
    THPVariable_Unpack(obj).unsafeGetTensorImpl()->set_python_dispatch(true);
  }
  return obj;
  END_HANDLE_TH_ERRORS
}

static PyObject* THPVariable_make_subclass(
    PyObject* _ignored,
    PyObject* args,
    PyObject* kwargs) {
  HANDLE_TH_ERRORS
  static PythonArgParser parser({
      "_make_subclass(PyObject* cls, Tensor data, bool require_grad=False, *, std::string_view? dispatch_sizes_strides_policy=None, bool dispatch_device=False, bool dispatch_layout=False, Device? device_for_backend_keys=None)",
  });
  ParsedArgs<7> parsed_args{};
  auto r = parser.parse(args, kwargs, parsed_args);
  PyObject* cls = r.pyobject(0);
  TORCH_CHECK_TENSOR_SUBTYPE(cls);
  // guard completely turns off torch dispatch modes, doesn't just pop off the
  // stack
  torch_dispatch_mode::StashTorchDispatchStackGuard td_g;
  c10::impl::DisablePythonDispatcher dpd_g;
  auto data =
      r.tensor(1).detach(); // creates a fresh Tensor (DEFINITELY_UNINITIALIZED)
  // We set `data`'s `allow_tensor_metadata_change` to true here, because we
  // want to allow the following use case for backward compatibility:
  //
  // ```python
  // rnn = torch.nn.RNN(100, 100, 2)
  // # The following calls `torch._cudnn_rnn_flatten_weight(rnn._flat_weights,
  // ...)`, # which changes storage of `rnn`'s weights in-place
  // rnn.flatten_parameters()
  // ```
  data.unsafeGetTensorImpl()->set_allow_tensor_metadata_change(true);
  data.set_requires_grad(r.toBool(2));
  const auto sizes_strides_policy = r.stringViewOptional(3);
  if (sizes_strides_policy.has_value()) {
    data.unsafeGetTensorImpl()->set_python_custom_sizes_strides(
        parseSizesStridesPolicyArgument(*sizes_strides_policy));
  }
  if (r.toBool(4)) {
    data.unsafeGetTensorImpl()->set_python_custom_device(true);
  }
  if (r.toBool(5)) {
    data.unsafeGetTensorImpl()->set_python_custom_layout(true);
  }
  if (!r.isNone(6)) {
    data.unsafeGetTensorImpl()->_change_backend_component_keys(r.device(6));
  }

  PyObject* obj = THPVariable_WrapWithType(data, (PyTypeObject*)cls);
  if (check_has_torch_dispatch(obj)) {
    THPVariable_Unpack(obj).unsafeGetTensorImpl()->set_python_dispatch(true);
  }
  return obj;
  END_HANDLE_TH_ERRORS
}

// Shared code factored out of THPVariable_make_wrapper_subclass and
// THPVariable_dtensor__new__.
static Tensor make_tensor_for_subclass_helper(
    SymIntArrayRef sym_sizes,
    OptionalSymIntArrayRef sym_strides,
    const std::optional<c10::SymInt>& sym_storage_offset,
    const TensorOptions& options,
    const std::optional<c10::SymInt>& storage_size,
    std::optional<DispatchKeySet> extra_dispatch_keys) {
  AutoDispatchBelowADInplaceOrView guard{}; // TODO: Remove.
  tracer::impl::NoTracerDispatchMode tracer_guard{};

  c10::SymInt size_bytes;
  auto dtype_itemsize = static_cast<int64_t>(options.dtype().itemsize());

  if (storage_size.has_value()) {
    size_bytes = storage_size.value();
  } else if (sym_strides.has_value()) {
    size_bytes = at::detail::computeStorageNbytes(
        sym_sizes,
        sym_strides.value(),
        dtype_itemsize,
        sym_storage_offset.value_or(0));
  } else {
    size_bytes = at::detail::computeStorageNbytesContiguous(
        sym_sizes, dtype_itemsize, sym_storage_offset.value_or(0));
  }

  // We use storages **only** to track aliasing of subclasses during tracing.
  // The actual data pointers are not valid.
  Storage storage{
      Storage::use_byte_size_t{},
      size_bytes,
      at::DataPtr{nullptr, options.device()},
      /*allocator=*/c10::GetAllocator(c10::kMeta),
      /*resizable=*/true};

  auto keys = c10::DispatchKeySet({options.computeDispatchKey()});
  if (extra_dispatch_keys.has_value()) {
    keys = keys | *extra_dispatch_keys;
  }
  Tensor tensor = at::detail::make_tensor<TensorImpl>(
      std::move(storage), keys, options.dtype());

  TensorImpl* tensor_impl = tensor.unsafeGetTensorImpl();

  if (sym_strides.has_value()) {
    tensor_impl->set_sizes_and_strides(
        sym_sizes, sym_strides.value(), sym_storage_offset);
  } else {
    TORCH_CHECK(
        !sym_storage_offset.has_value(),
        "setting storage offset without stride not supported");
    tensor_impl->generic_set_sizes_contiguous(sym_sizes);
  }
  return tensor;
}

static PyObject* THPVariable_make_wrapper_subclass(
    PyObject* /*unused*/,
    PyObject* args,
    PyObject* kwargs) {
  HANDLE_TH_ERRORS
  // NB: pin_memory doesn't actually do anything
  // TODO: strides variant?

  // cls: Python subclass type
  // size, strides, storage_offset, memory_format, dtype: self-explanatory
  // layout: memory layout, e.g. for types of Nested Tensors or other sparse
  //         tensors
  // pin_memory, requires_grad: self-explanatory
  // dispatch_sizes_strides_policy: string - which sizes/strides we should
  //                                dispatch to a custom python implementation.
  // dispatch_device: whether to dispatch to a custom python implementation
  //                  for device
  // dispatch_layout: whether to dispatch to a custom python implementation
  //                  for layout
  // _extra_dispatch_keys: additional dispatch keys to add to the tensor
  // storage_size: if provided, skip storage size calculation and just use the
  //               value provided. One use case is for Nested Tensor, where the
  //               storage size cannot be calculated from the sizes/strides
  //               (because they contain a NestedInt).
  static PythonArgParser parser({
      "_make_wrapper_subclass(PyObject* cls, SymIntArrayRef size, SymIntArrayRef? strides=None, "
      "SymInt? storage_offset=None, MemoryFormat? memory_format=None, ScalarType dtype=None, "
      "Layout layout=torch.strided, Device device=None, bool pin_memory=False, bool requires_grad=False, "
      "std::string_view? dispatch_sizes_strides_policy=None, bool dispatch_device=False, bool dispatch_layout=False, "
      "DispatchKeySet _extra_dispatch_keys=None, SymInt? storage_size=None)",
  });
  ParsedArgs<15> parsed_args{};
  auto r = parser.parse(args, kwargs, parsed_args);
  PyObject* cls = r.pyobject(0);

  TORCH_CHECK_TENSOR_SUBTYPE(cls);

  // This is an important safety check; without it, the default behavior will be
  // to continue on to the underlying CPU/CUDA kernel advertised by the dispatch
  // key, which will immediately segfault because the data pointer is null.  By
  // forcing users to define __torch_dispatch__ we ensure this does not happen
  // TODO: This check is not complete; because the user can disable torch
  // dispatch and then go again, triggering segfault.  TBH I'm thinking I want
  // to delete this function entirely
  py::object attr = PyObject_FastGetAttrString(cls, "__torch_dispatch__");
  TORCH_CHECK_TYPE(
      attr.ptr() != nullptr &&
          attr.ptr() != torch::disabled_torch_dispatch_impl(),
      ((PyTypeObject*)cls)->tp_name,
      " must define __torch_dispatch__");

  const auto options = TensorOptions()
                           .dtype(r.scalartype(5))
                           .device(r.device(7))
                           .layout(r.layoutOptional(6))
                           // NB: long standing issue, requires_grad is not
                           // respected here; you have to set it post facto, see
                           // https://github.com/pytorch/pytorch/issues/26428
                           // .requires_grad(r.toBool(7))
                           .pinned_memory(r.toBool(8));

  // don't bother releasing GIL here, as we are not allocating any nontrivial
  // data
  auto sym_sizes = r.symintlist(1);
  auto sym_strides_own = r.symintlistOptional(2);
  Tensor tensor = make_tensor_for_subclass_helper(
      /*sym_sizes=*/r.symintlist(1),
      /*sym_strides=*/r.symintlistOptional(2),
      /*sym_storage_offset=*/r.toSymIntOptional(3),
      options,
      /*storage_size=*/r.toSymIntOptional(14),
      r.toDispatchKeySetOptional(13));

  tensor.unsafeGetTensorImpl()->set_python_dispatch(true);

  const auto sizes_strides_policy = r.stringViewOptional(10);
  if (sizes_strides_policy.has_value()) {
    tensor.unsafeGetTensorImpl()->set_python_custom_sizes_strides(
        parseSizesStridesPolicyArgument(*sizes_strides_policy));
  }

  tensor.set_requires_grad(r.toBool(9));

  if (r.toBool(11)) {
    tensor.unsafeGetTensorImpl()->set_python_custom_device(true);
  }
  if (r.toBool(12)) {
    tensor.unsafeGetTensorImpl()->set_python_custom_layout(true);
  }

  return THPVariable_WrapWithType(std::move(tensor), (PyTypeObject*)cls);
  END_HANDLE_TH_ERRORS
}

#if IS_PYBIND_2_13_PLUS
#define DEFINE_CACHING_PYTHON_IMPORT_GETTER(name, import_expr)             \
  static py::handle name() {                                               \
    PYBIND11_CONSTINIT static py::gil_safe_call_once_and_store<py::object> \
        storage;                                                           \
    return storage                                                         \
        .call_once_and_store_result(                                       \
            []() -> py::object { return import_expr; })                    \
        .get_stored();                                                     \
  }
#else
#define DEFINE_CACHING_PYTHON_IMPORT_GETTER(name, import_expr)     \
  static py::handle name() {                                       \
    static py::handle storage = py::object(import_expr).release(); \
    return storage;                                                \
  }
#endif

DEFINE_CACHING_PYTHON_IMPORT_GETTER(
    get_dtensor_class_impl,
    py::module::import("torch.distributed.tensor").attr("DTensor"))

py::handle get_dtensor_class() {
  return get_dtensor_class_impl();
}

DEFINE_CACHING_PYTHON_IMPORT_GETTER(
    get_dtensor_spec_class,
    py::module::import("torch.distributed.tensor")
        .attr("_dtensor_spec")
        .attr("DTensorSpec"))

DEFINE_CACHING_PYTHON_IMPORT_GETTER(
    get_replicate_class,
    py::module::import("torch.distributed.tensor")
        .attr("placement_types")
        .attr("Replicate"))

DEFINE_CACHING_PYTHON_IMPORT_GETTER(
    get_tensor_meta_class,
    py::module::import("torch.distributed.tensor")
        .attr("_dtensor_spec")
        .attr("TensorMeta"))

DEFINE_CACHING_PYTHON_IMPORT_GETTER(
    get_dtensor_op_dispatcher,
    py::module::import("torch.distributed.tensor")
        .attr("DTensor")
        .attr("_op_dispatcher"))

DEFINE_CACHING_PYTHON_IMPORT_GETTER(
    get_dtensor_dispatch,
    py::module::import("torch.distributed.tensor")
        .attr("DTensor")
        .attr("_op_dispatcher")
        .attr("_dispatch_fast_path_python_tail"))

DEFINE_CACHING_PYTHON_IMPORT_GETTER(
    get_dtensor_dispatcher_wrap,
    py::module::import("torch.distributed.tensor")
        .attr("DTensor")
        .attr("_op_dispatcher")
        .attr("wrap"))

DEFINE_CACHING_PYTHON_IMPORT_GETTER(
    get_dtensor_get_local_results_slow_path,
    py::module::import("torch")
        .attr("distributed")
        .attr("tensor")
        .attr("DTensor")
        .attr("_op_dispatcher")
        .attr("_dispatch_get_local_results_slow_path"))

DEFINE_CACHING_PYTHON_IMPORT_GETTER(
    get_output_sharding_class,
    py::module::import("torch.distributed.tensor")
        .attr("_op_schema")
        .attr("OutputSharding"))

DEFINE_CACHING_PYTHON_IMPORT_GETTER(
    get_op_strategy_class,
    py::module::import("torch.distributed.tensor")
        .attr("_op_schema")
        .attr("OpStrategy"))

DEFINE_CACHING_PYTHON_IMPORT_GETTER(
    get_tuple_strategy_class,
    py::module::import("torch.distributed.tensor")
        .attr("_op_schema")
        .attr("TupleStrategy"))

static bool arg_type_tensor_or_tensor_list_like(py::handle arg) {
  const auto dtensor_spec_class = get_dtensor_spec_class();
  const auto op_strategy_class = get_op_strategy_class();
  const auto tuple_strategy_class = get_tuple_strategy_class();

  if (py::isinstance(arg, dtensor_spec_class) ||
      py::isinstance(arg, op_strategy_class) ||
      py::isinstance(arg, tuple_strategy_class)) {
    return true;
  }
  if (!PyList_Check(arg.ptr())) {
    return false;
  }
  py::list arg_list = py::reinterpret_borrow<py::list>(arg);
  for (const auto e : arg_list) {
    if (!e.is_none() && !py::isinstance(e, dtensor_spec_class) &&
        !py::isinstance(e, op_strategy_class) &&
        !py::isinstance(e, tuple_strategy_class)) {
      return false;
    }
  }
  return true;
}

#if IS_PYTHON_3_11_PLUS
#define MAYBE_FOR_EACH_PYTHON_3_10_MINUS_DTENSOR_INTERNED_STRING(_)
#else
#define MAYBE_FOR_EACH_PYTHON_3_10_MINUS_DTENSOR_INTERNED_STRING(_) _(__name__)
#endif

#define FOR_EACH_DTENSOR_INTERNED_STRING(_)                   \
  MAYBE_FOR_EACH_PYTHON_3_10_MINUS_DTENSOR_INTERNED_STRING(_) \
  _(_comparison_key)                                          \
  _(_custom_op_handlers)                                      \
  _(_local_tensor)                                            \
  _(_spec)                                                    \
  _(_unwrap_to_op_info_impl)                                  \
  _(args_schema)                                              \
  _(compute_mesh)                                             \
  _(device_mesh)                                              \
  _(dtype)                                                    \
  _(get_coordinate)                                           \
  _(kwargs_schema)                                            \
  _(ndim)                                                     \
  _(needs_pytree)                                             \
  _(needs_redistribute)                                       \
  _(op)                                                       \
  _(op_to_schema_info)                                        \
  _(op_to_schema_info_for_single_dim_strategy)                \
  _(output_sharding)                                          \
  _(output_spec)                                              \
  _(schema_info)                                              \
  _(shape)                                                    \
  _(sharding_propagator)                                      \
  _(size)                                                     \
  _(static_argnum)                                            \
  _(static_kwargkey)                                          \
  _(stride)                                                   \
  _(tensor_meta)

struct DTensorInternedStrings {
#define DECLARE_INTERNED_STRING_VARIABLE(s) PyObject* s;
  FOR_EACH_DTENSOR_INTERNED_STRING(DECLARE_INTERNED_STRING_VARIABLE)
#undef DECLARE_INTERNED_STRING_VARIABLE
};

static DTensorInternedStrings dtensor_interned_strings;

#ifdef USE_DISTRIBUTED
static bool intern_dtensor_strings() {
#define INTERN_DTENSOR_STRING(s)                                           \
  TORCH_INTERNAL_ASSERT_DEBUG_ONLY(dtensor_interned_strings.s == nullptr); \
  dtensor_interned_strings.s = PyUnicode_InternFromString(#s);             \
  if (dtensor_interned_strings.s == nullptr) {                             \
    return false;                                                          \
  }

  FOR_EACH_DTENSOR_INTERNED_STRING(INTERN_DTENSOR_STRING);
#undef INTERN_DTENSOR_STRING
  return true;
}
#endif

static bool checked_not(PyObject* obj) {
  int result = PyObject_Not(obj);
  if (result == -1) {
    throw py::error_already_set();
  }
  return result;
}

static bool checked_istrue(PyObject* obj) {
  int result = PyObject_IsTrue(obj);
  if (result == -1) {
    throw py::error_already_set();
  }
  return result;
}

// pybind11 does not not use PyObject_Vectorcall currently; it seems
// to materialize a tuple of args instead.
template <std::size_t N>
static py::object checked_vectorcall(
    PyObject* obj,
    std::array<PyObject*, N> args) {
  PyObject* result = PyObject_Vectorcall(obj, args.data(), N, nullptr);
  if (!result) {
    throw py::error_already_set();
  }
  return py::reinterpret_steal<py::object>(result);
}

template <typename... Args>
static py::object checked_vectorcall(PyObject* obj, Args... args) {
  static_assert(
      (std::is_same_v<Args, PyObject*> && ...),
      "must pass PyObject* to checked_vectorcall!");
  std::array<PyObject*, sizeof...(Args)> arr = {args...};
  return checked_vectorcall(obj, arr);
}

static c10::SymDimVector tuple_to_symintlist(PyObject* obj) {
  TORCH_INTERNAL_ASSERT_DEBUG_ONLY(PyTuple_Check(obj));
  c10::SymDimVector res;
  const auto size = PyTuple_GET_SIZE(obj);
  res.reserve(size);
  for (const auto idx : c10::irange(size)) {
    PyObject* item = PyTuple_GET_ITEM(obj, idx);
    if (THPUtils_checkLongExact(item)) {
      res.emplace_back(THPUtils_unpackLong(item));
    } else if (torch::is_symint(py::handle(item))) {
      res.push_back(py::handle(item).cast<c10::SymInt>());
    } else {
      // N.B. torch.Tensor.__index__ exists, so this should handle
      // scalar Tensors fine.
      res.emplace_back(THPUtils_unpackIndex(item));
    }
  }
  return res;
}

// As a Python object, DTensorSpec can be stored directly within
// IValue, but doing so is inefficient -- it requires a
// heap-allocated, reference counted intermediate
// ivalue::PyObjectHolder.
// Representation options:
// 1) Add an IValue tag to represent a placeholder object.
// 2) Play representational tricks -- stuff information into an IValue
// payload, such as by creating impossible
// intrusive_ptr_target*. Problem: this would cause IValue copying and
// possibly destruction to crash and so would be horribly unsafe.
// 3) Represent DTensorSpec directly inside IValue despite the inefficiency.
// 4) Leave the actual DTensor in the list of IValues, but detect it efficiently
// and transparently replace.
// 5) Just use a 24-byte struct of IValue + extra py::object.
//
// Given the high blast radius of (1), the unsafety of (2), the likely
// poor performance of (3), and detection of (4) looking less
// efficient than (5), (5) seems like the best path forward.

// We can't safely steal bits from IValue, so we just use 24 bytes of
// space. If dtensor_spec is non-null (truthy) then it's the active
// member, otherwise it's iv.
struct IValueOrDTensorSpec {
  IValueOrDTensorSpec() = default;
  explicit IValueOrDTensorSpec(c10::IValue v) : iv(std::move(v)) {}
  explicit IValueOrDTensorSpec(py::object dts) : dtensor_spec(std::move(dts)) {}
  c10::IValue iv;
  py::object dtensor_spec;

  bool operator==(const IValueOrDTensorSpec& rhs) const {
    return dtensor_spec
        ? (rhs.dtensor_spec && dtensor_spec.equal(rhs.dtensor_spec))
        : (iv == rhs.iv);
  }
};

// This corresponds to the Python OpSchema class in that it is the key
// for the (native version of the) sharding propagator cache. It is
// missing essentially everything else from the Python OpSchema
// though.
class NativeOpSchema {
 public:
  NativeOpSchema(
      const c10::OperatorHandle& op,
      c10::SmallVector<IValueOrDTensorSpec, 8> comparison_key,
      std::size_t comparison_key_hash,
      std::size_t args_schema_len)
      : op_(op),
        hash_(hash_combine(
            hash_combine(
                std::hash<c10::OperatorHandle>()(op),
                comparison_key_hash),
            args_schema_len)),
        args_schema_len_(args_schema_len),
        comparison_key_(std::move(comparison_key)) {}

  bool operator==(const NativeOpSchema& rhs) const {
    // If two NativeOpSchema are being compared, they are probably
    // equal, because comparison is occurring during a hash table
    // lookup and we know the hashes are already equal. Therefore, we
    // don't bother checking hash_ first.
    return op_ == rhs.op_ && args_schema_len_ == rhs.args_schema_len_ &&
        comparison_key_ == rhs.comparison_key_;
  }

  std::size_t hash() const noexcept {
    return hash_;
  }

  const c10::OperatorHandle& op() const {
    return op_;
  }

  const c10::SmallVector<IValueOrDTensorSpec, 8>& comparison_key() const {
    return comparison_key_;
  }

  // Format schema as string for logging
  std::string format_inputs() const {
    std::ostringstream ss;
    ss << op_.operator_name().name;
    if (!op_.operator_name().overload_name.empty()) {
      ss << "." << op_.operator_name().overload_name;
    }
    ss << "(";
    bool first = true;
    for (const auto& item : comparison_key_) {
      if (!first)
        ss << ", ";
      first = false;
      if (item.dtensor_spec) {
        ss << py::str(item.dtensor_spec).cast<std::string>();
      } else if (!item.iv.isNone()) {
        ss << item.iv;
      }
    }
    ss << ")";
    return ss.str();
  }

 private:
  // It would *not* be correct to store this by reference, because we
  // have no guarantees about its lifetime. This class is cheap anyway.
  c10::OperatorHandle op_;
  std::size_t hash_;
  // Subtle point: consider clamp.Tensor(Tensor self, Tensor?
  // min=None, Tensor? max=None). The invocations clamp(t1, None, t2)
  // and clamp(t1, t2, None) have the same comparison key (t1, t2)
  // because we drop non-static non-tensor args from comparison. The
  // only way we happen to be able to tell them apart is that we omit
  // trailing defaulted arguments from the args tuple passed to
  // __torch_dispatch__ (and hence to DTensor dispatch as well), so
  // they have different args_schema_len_.
  //
  // I am preserving this existing behavior, but I suspect we should
  // make an algorithm change to be less brittle, such as including
  // None defaults for Tensor arguments in the comparison.
  std::size_t args_schema_len_;
  // There is no particular justification for the choice of 8
  // here. Feel free to change it.
  c10::SmallVector<IValueOrDTensorSpec, 8> comparison_key_;
};

namespace std {
template <>
struct hash<NativeOpSchema> {
  std::size_t operator()(const NativeOpSchema& schema) const noexcept {
    return schema.hash();
  }
};
} // namespace std

// Helper to check if dtensor debug logging is enabled and log cache hits.
// We cache the logger and debug-enabled check to avoid repeated Python calls.
// Pattern follows torch/csrc/jit/passes/onnx/onnx_log.cpp
namespace {
thread_local py::object dtensor_dispatch_logger;
thread_local py::object logging_DEBUG;
thread_local bool dtensor_dispatch_logger_initialized = false;
thread_local bool dtensor_debug_logging_enabled = false;

void init_dtensor_dispatch_logger() {
  if (!dtensor_dispatch_logger_initialized) {
    dtensor_dispatch_logger_initialized = true;
    try {
      auto logging = py::module_::import("logging");
      logging_DEBUG = logging.attr("DEBUG");
      dtensor_dispatch_logger =
          logging.attr("getLogger")("torch.distributed.tensor._dispatch");
      dtensor_debug_logging_enabled =
          dtensor_dispatch_logger.attr("isEnabledFor")(logging_DEBUG)
              .cast<bool>();
    } catch (...) {
      dtensor_dispatch_logger = py::none();
      dtensor_debug_logging_enabled = false;
    }
  }
}

void log_sharding_prop_cache_hit(
    const NativeOpSchema& schema,
    const py::object& cached_sharding) {
  init_dtensor_dispatch_logger();
  if (!dtensor_debug_logging_enabled) {
    return;
  }
  std::ostringstream ss;
  ss << "sharding_prop HIT (C++ fast path): ";
  ss << schema.format_inputs();
  auto output_spec = cached_sharding.attr("output_spec");
  if (!output_spec.is_none()) {
    ss << " -> " << py::str(output_spec).cast<std::string>();
  }
  dtensor_dispatch_logger.attr("debug")(ss.str());
}
} // namespace

// Map from OpSchema to pyobject sharding propagation config.
class NativeShardingPropagatorCache {
 public:
  // Returns an invalid (falsey) py::object if the lookup fails.
  py::object find(const NativeOpSchema& op_schema) const {
    if (auto it = repr_.find(op_schema); it != repr_.end()) {
      hits_++;
      return py::object(it->second);
    }
    misses_++;
    return py::object();
  }

  void insert(NativeOpSchema&& op_schema, py::object output_sharding) {
    auto [it, inserted] =
        repr_.emplace(std::move(op_schema), std::move(output_sharding));
    TORCH_INTERNAL_ASSERT(
        inserted,
        "tried to insert already-present element in NativeShardingPropagatorCache!");
  }

  auto hits() const {
    return hits_;
  }

  auto misses() const {
    return misses_;
  }

 private:
  c10::FastMap<NativeOpSchema, py::object> repr_;
  // Cache is thread-local, so we don't take any further action for
  // thread-safety of these.
  mutable std::size_t hits_ = 0;
  mutable std::size_t misses_ = 0;
};

static std::optional<std::pair<NativeOpSchema, /*ComputeMesh*/ py::object>>
create_native_op_schema(
    const c10::OperatorHandle& op,
    py::handle py_op,
    torch::jit::Stack* stack);

static std::mutex native_sharding_propagator_cache_cleanup_mutex;
static c10::
    FastMap<std::thread::id, std::optional<NativeShardingPropagatorCache>*>
        all_thread_caches;
thread_local std::optional<NativeShardingPropagatorCache>
    native_sharding_propagator_cache_DO_NOT_USE;

NativeShardingPropagatorCache&
get_thread_local_native_sharding_propagator_cache() {
  if (!native_sharding_propagator_cache_DO_NOT_USE.has_value()) {
    native_sharding_propagator_cache_DO_NOT_USE.emplace();
    std::lock_guard<std::mutex> lock(
        native_sharding_propagator_cache_cleanup_mutex);
    const auto this_thread_id = std::this_thread::get_id();
    all_thread_caches[this_thread_id] =
        &native_sharding_propagator_cache_DO_NOT_USE;
    py::dict thread_dict =
        py::reinterpret_borrow<py::dict>(PyThreadState_GetDict());
    // We need to clean up before Python detaches from the thread if
    // the thread is being destroyed.
    if (!thread_dict.contains("__DTensor_fastpath_thread_cache_cleanup")) {
      thread_dict["__DTensor_fastpath_thread_cache_cleanup"] =
          py::capsule(new std::thread::id(this_thread_id), [](void* p) {
            auto* ptid = reinterpret_cast<std::thread::id*>(p);
            {
              std::lock_guard<std::mutex> inner_lock(
                  native_sharding_propagator_cache_cleanup_mutex);
              auto it = all_thread_caches.find(*ptid);
              if (it != all_thread_caches.end()) {
                // We need to both:
                // 1) free python objects, and
                it->second->reset();
                // 2) make sure we don't try to come back and mess with
                // a destroyed thread-local at module unload (e.g.,
                // process exit) time.
                all_thread_caches.erase(it);
              }
            }
            delete ptid;
          });
    }
  }
  return native_sharding_propagator_cache_DO_NOT_USE.value();
}

// We need to clean up all thread_locals if our module is getting
// unloaded.
void cleanup_thread_local_native_sharding_propagator_caches() {
  std::lock_guard<std::mutex> lock(
      native_sharding_propagator_cache_cleanup_mutex);
  for (auto& [_, popt_cache] : all_thread_caches) {
    popt_cache->reset();
  }
  all_thread_caches.clear();
}

static void replace_dtensors_with_local_tensor(torch::jit::Stack& stack);

static bool is_default_overload(const std::string& overload_name) {
  return overload_name.empty() || overload_name == "default";
}

static bool is_random_op(const c10::OperatorHandle& op) {
  // NOTE: must stay in sync with _random_ops in
  // torch/distributed/tensor/_dispatch.py
  constexpr auto aten_namespace_prefix_len = 6;
  const auto& op_name = op.operator_name();
  if (op_name.name.size() <= aten_namespace_prefix_len ||
      memcmp(op_name.name.data(), "aten::", aten_namespace_prefix_len) != 0) {
    return false;
  }
  static constexpr std::array<std::string_view, 6> random_names = {{
      "native_dropout",
      "normal_",
      "rand_like",
      "randn_like",
      "uniform_",
      "bernoulli",
  }};
  std::string_view name_without_namespace(
      op_name.name.c_str() + aten_namespace_prefix_len,
      op_name.name.size() - aten_namespace_prefix_len);
  if (name_without_namespace == "bernoulli_") {
    return op_name.overload_name == "float";
  }
  if (name_without_namespace == "randint_like") {
    return is_default_overload(op_name.overload_name) ||
        op_name.overload_name == "low_dtype" ||
        op_name.overload_name == "low_dtype_out";
  }
  const auto it = std::find(
      random_names.begin(), random_names.end(), name_without_namespace);
  if (it == random_names.end()) {
    return false;
  }
  return is_default_overload(op_name.overload_name);
}

// Puts local results on the stack. Return true for success, false for bailout
// to slow path.
static bool get_local_results(
    const c10::OperatorHandle& op,
    py::handle output_sharding,
    py::handle compute_mesh,
    bool participating,
    torch::jit::Stack* stack) {
  if (participating) {
    // computation that happens in the current rank of the mesh, normal case
    if (checked_istrue(
            output_sharding.attr(dtensor_interned_strings.needs_redistribute)
                .ptr()) ||
        is_random_op(op)) {
      // Bail out to slow path.
      return false;
    }
    // normal case, run local sharded op computation.

    // It is slightly inefficient that we take another pass over
    // arguments here when we just did one in create_native_op_schema to
    // create the comparison key. However, we have a crucial difference:
    // in the NativeOpSchema, we don't want to waste time dealing with
    // defaulted args. Here, we need to provide defaulted args because
    // we are going to make a local op call.
    replace_dtensors_with_local_tensor(*stack);
    op.callBoxed(*stack);
  } else {
    // For a non-participating device (happens on rank that does not
    // belong to the device mesh), we do:
    //
    //   1. if the return type is scalar, set the local result to
    //   None.
    //   2. if the return type is Tensor or List[Tensor], return
    //   empty tensor(s) with correct dtype.

    stack->clear();

    auto spec = output_sharding.attr(dtensor_interned_strings.output_spec);
    if (spec.is_none()) {
      // For a scalar return type, the non-participating device has
      // None as its local result.
      stack->emplace_back(); // Return None.
      return true;
    }

    const auto default_tensor = [](py::handle spec) -> Tensor {
      auto tensor_meta = spec.attr(dtensor_interned_strings.tensor_meta);
      TORCH_CHECK(
          !tensor_meta.is_none(), py::str(spec), " has no tensor metadata.");
      const auto sizes = tensor_meta.attr(dtensor_interned_strings.shape);
      TORCH_CHECK(
          PyTuple_Check(sizes.ptr()), "spec.tensor_meta.shape must be a tuple");
      const auto dtype = tensor_meta.attr(dtensor_interned_strings.dtype);
      TORCH_CHECK(
          THPDtype_Check(dtype.ptr()),
          "spec.tensor_meta.dtype must be a torch.dtype");
      const auto scalar_type =
          reinterpret_cast<THPDtype*>(dtype.ptr())->scalar_type;
      if (py::cast<py::tuple>(sizes).empty()) {
        // scalar tensor
        return at::zeros({}, scalar_type);
      } else {
        // non-scalar tensor
        return at::empty({0}, scalar_type);
      }
    };
    auto handle_sequence = [&default_tensor, &op, stack](auto sequence) {
      c10::List<c10::IValue> result(op.schema().returns().at(0).type());
      for (const auto& item : sequence) {
        TORCH_CHECK(
            !item.is_none(),
            "return type ",
            op.schema().returns().at(0).type(),
            " in DTensor op is not supported");
        result.push_back(default_tensor(item));
      }
      stack->emplace_back(std::move(result));
    };

    if (py::isinstance(spec, get_dtensor_spec_class())) {
      stack->emplace_back(default_tensor(spec));
    } else if (PyList_Check(spec.ptr())) {
      handle_sequence(py::reinterpret_borrow<py::list>(spec));
    } else if (PyTuple_Check(spec.ptr())) {
      handle_sequence(py::reinterpret_borrow<py::tuple>(spec));
    } else if (PySequence_Check(spec.ptr())) {
      handle_sequence(py::reinterpret_borrow<py::sequence>(spec));
    } else {
      // return None.
      stack->emplace_back();
    }
  }
  return true;
}

static void functionalize_unsafe_set(at::Tensor& dst, const at::Tensor& src) {
  at::native::checkSetStorage(
      dst,
      src.storage(),
      dst.sym_storage_offset(),
      dst.sym_sizes(),
      dst.sym_strides(),
      /*check_offset_in_bounds=*/false);
}

static bool sets_intersect(
    const std::unordered_set<Symbol>& smaller,
    const std::unordered_set<Symbol>& bigger) {
  if (smaller.size() > bigger.size()) {
    return sets_intersect(bigger, smaller);
  }
  for (const auto& item : smaller) {
    if (bigger.find(item) != bigger.end()) {
      return true;
    }
  }
  return false;
}

py::object dispatchDTensorOp(
    const c10::OperatorHandle& op,
    py::handle py_op,
    py::handle args,
    py::handle kwargs,
    torch::jit::Stack* stack) {
  py::object cached_sharding;
  const auto op_dispatcher = get_dtensor_op_dispatcher();
  {
    const auto custom_op_handlers =
        op_dispatcher.attr(dtensor_interned_strings._custom_op_handlers);
    TORCH_CHECK(
        PyDict_Check(custom_op_handlers.ptr()),
        "_custom_op_handlers must be a dict!");
    PyObject* custom_op_handler =
        PyDict_GetItemWithError(custom_op_handlers.ptr(), py_op.ptr());
    if (custom_op_handler) {
      auto result = checked_vectorcall(
          custom_op_handler, py_op.ptr(), args.ptr(), kwargs.ptr());
      stack->clear();
      return result;
    } else if (PyErr_Occurred()) {
      throw py::error_already_set();
    }
  }

  torch::jit::Stack saved_args = *stack;
  NativeShardingPropagatorCache* native_sharding_propagator_cache = nullptr;
  // In the original Python implementation of DTensor dispatch, the creation
  // of OpInfo (which includes the OpSchema computed here) never fails. However,
  // C++ support for all the features of OpSchema are not supported; in this
  // case opt_native_op_schema is nullopt.  In this case, we need to fallback
  // to the Python logic for doing so.  If you are comparing against the old
  // Python code, this is a bit tricky, since the Python 'dispatch' function
  // has been completely deleted.

  // First, we will try to short-circuit Python entirely using the fast path.
  // Here, we never materialize OpInfo, we generate a gimped NativeOpSchema
  // object which has exactly the information you need to do a hash lookup.
  auto opt_native_op_schema = create_native_op_schema(op, py_op, stack);
  if (opt_native_op_schema.has_value()) {
    native_sharding_propagator_cache =
        &get_thread_local_native_sharding_propagator_cache();
    cached_sharding =
        native_sharding_propagator_cache->find(opt_native_op_schema->first);
  }

  py::object py_op_info;
  if (!cached_sharding) {
    // OK, the C++ fastpath failed.  Let's use the Python path to generate the
    // OpInfo (which is guaranteed to work), which we will need to either
    // redo the cache lookup or compute the value for real.
    py_op_info = checked_vectorcall(
        op_dispatcher.attr("unwrap_to_op_info").ptr(),
        py_op.ptr(),
        args.ptr(),
        kwargs.ptr());

    py::object sharding = checked_vectorcall(
        op_dispatcher.attr("_propagate_op_sharding_dispatch_slow_path").ptr(),
        py_op.ptr(),
        args.ptr(),
        kwargs.ptr(),
        py_op_info.ptr(),
        /*try_cache*/ !opt_native_op_schema.has_value() ? Py_True : Py_False);
    // This is a hack, because the dispatch slow path sometimes returns
    // a sharding result (in which case we need to keep going) but it
    // will sometimes just decompose and directly return a Tensor result,
    // in which case we should return immediately.  In this case, sharding
    // is not a sharding at all; it's the real result!
    if (!py::isinstance(sharding, get_output_sharding_class())) {
      stack->clear();
      return sharding;
    }
    cached_sharding = sharding;
    if (opt_native_op_schema.has_value()) {
      native_sharding_propagator_cache->insert(
          std::move(opt_native_op_schema->first), std::move(sharding));
    }
    py_op_info.attr(dtensor_interned_strings.output_sharding) = cached_sharding;
  } else if (opt_native_op_schema.has_value()) {
    // Cache hit - log it if debug logging is enabled
    log_sharding_prop_cache_hit(opt_native_op_schema->first, cached_sharding);
  }

  const auto get_py_op_info_if_needed = [&, &args = args, &kwargs = kwargs]() {
    if (!py_op_info) {
      py_op_info = checked_vectorcall(
          op_dispatcher.attr(dtensor_interned_strings._unwrap_to_op_info_impl)
              .ptr(),
          py_op.ptr(),
          args.ptr(),
          kwargs.ptr(),
          Py_False);
      py_op_info.attr(dtensor_interned_strings.output_sharding) =
          cached_sharding;
    }
  };

  TORCH_INTERNAL_ASSERT_DEBUG_ONLY(
      !kwargs.is_none(),
      "Python op_dispatch implementation expects non-None kwargs");

  py::object compute_mesh;
  if (opt_native_op_schema.has_value()) {
    compute_mesh = std::move(opt_native_op_schema->second);
  } else {
    get_py_op_info_if_needed();
    compute_mesh = py_op_info.attr(dtensor_interned_strings.compute_mesh);
  }

  const bool participating =
      !checked_vectorcall(
           compute_mesh.attr(dtensor_interned_strings.get_coordinate).ptr())
           .is_none();
  const bool local_results_success = get_local_results(
      op, cached_sharding, compute_mesh, participating, stack);
  py::object py_local_results;
  if (local_results_success) {
    py_local_results = torch::jit::createPyObjectForStack(std::move(*stack));
  } else {
    get_py_op_info_if_needed();
    py_local_results = checked_vectorcall(
        get_dtensor_get_local_results_slow_path().ptr(),
        py_op.ptr(),
        args.ptr(),
        py_op_info.ptr());
  }

  const auto& operator_name = op.operator_name();
  // Simple analysis of function schema to determine if this is an
  // inplace variant. It might not be entirely correct, but it's good
  // enough for now.
  const bool is_inplace_op =
      !operator_name.name.empty() && operator_name.name.back() == '_';
  // Simple analysis of function schema to determine if this is an
  // ou variant. It might not be entirely correct, but it's good
  // enough for now.
  const bool is_out_variant_op = !is_inplace_op &&
      operator_name.overload_name.find("out") != std::string::npos;

  // Fast path for default or view ops.
  const auto output_spec =
      cached_sharding.attr(dtensor_interned_strings.output_spec);
  if (!is_inplace_op && !is_out_variant_op &&
      !(output_spec.is_none() &&
        (op.operator_name().name == "aten::equal" &&
         is_default_overload(op.operator_name().overload_name)))) {
    const auto wrap = get_dtensor_dispatcher_wrap();
    auto wrapped_result = checked_vectorcall(
        wrap.ptr(), py_local_results.ptr(), output_spec.ptr());
    if (!participating) {
      stack->clear();
      return wrapped_result;
    }

    // Direct C++ implementation of return_and_correct_aliasing for view ops.

    // py::tuple's default constructor allocates a size-0 tuple, so we
    // wrap in optional to get a detectable empty state.
    std::optional<py::tuple> wrapped_result_tuple;
    if (PyTuple_Check(wrapped_result.ptr())) {
      wrapped_result_tuple = py::reinterpret_borrow<py::tuple>(wrapped_result);
    }
    const auto& returns = op.schema().returns();
    const auto num_arguments = op.schema().arguments().size();
    for (const auto arg_idx : c10::irange(num_arguments)) {
      const auto& arg_schema = op.schema().arguments()[arg_idx];
      const auto* arg_alias_info = arg_schema.alias_info();
      if (!arg_alias_info || arg_alias_info->isWrite()) {
        continue;
      }
      // If we ever get here, it's a view op. Therefore, it does not
      // have mutable output aliases, so we skip that portion of
      // return_and_correct_aliasing. Furthermore, we *only* want to
      // return_and_correct_aliasing if it's a view op, so we do not
      // need to port the mutable output aliases portion of
      // return_and_correct_aliasing at all.
      const c10::IValue& arg_iv =
          saved_args.at(saved_args.size() - num_arguments + arg_idx);
      if (!arg_iv.isTensor()) {
        continue;
      }
      const auto& arg = arg_iv.toTensor();
      int ret_idx = 0;
      for (const auto& ret_schema : returns) {
        const auto* ret_alias_info = ret_schema.alias_info();
        if (!ret_alias_info) {
          ret_idx++;
          continue;
        }
        if (sets_intersect(
                arg_alias_info->beforeSets(), ret_alias_info->beforeSets())) {
          py::object ret;
          if (wrapped_result_tuple.has_value()) {
            ret = wrapped_result_tuple.value()[ret_idx];
          } else {
            TORCH_INTERNAL_ASSERT(ret_idx == 0);
            ret = wrapped_result;
          }
          if (PyList_Check(ret.ptr())) {
            py::list ret_list = py::reinterpret_borrow<py::list>(ret);
            for (const auto& r : ret_list) {
              auto tensor = py::cast<at::Tensor>(r);
              functionalize_unsafe_set(tensor, arg);
            }
          } else {
            auto tensor = py::cast<at::Tensor>(ret);
            functionalize_unsafe_set(tensor, arg);
          }
        }
        ret_idx++;
      }
    }
    stack->clear();
    return wrapped_result;
  }

  auto dispatch = get_dtensor_dispatch();
  auto result = checked_vectorcall(
      dispatch.ptr(),
      py_op.ptr(),
      args.ptr(),
      kwargs.ptr(),
      compute_mesh.ptr(),
      cached_sharding.ptr(),
      py_local_results.ptr(),
      participating ? Py_True : Py_False,
      is_inplace_op ? Py_True : Py_False,
      is_out_variant_op ? Py_True : Py_False);
  stack->clear();
  return result;
}

// DTensor-specific variant of make_wrapper_subclass to minimize DTensor
// overhead.
static PyObject* THPVariable_dtensor_new(
    PyObject* /*unused*/,
    PyObject* args,
    PyObject* kwargs) {
  HANDLE_TH_ERRORS
  static PythonArgParser parser({
      "_dtensor__new__(PyObject* cls, Tensor local_tensor, PyObject* spec, bool requires_grad)",
  });
  ParsedArgs<4> parsed_args{};
  auto r = parser.parse(args, kwargs, parsed_args);
  PyObject* cls = r.pyobject(0);

  TORCH_CHECK_TENSOR_SUBTYPE(cls);

#ifndef NDEBUG
  // This is specifically for making a DTensor, which we know defines
  // __torch_dispatch__. Check anyway in debug builds in case somebody
  // removes it.
  py::object attr = PyObject_FastGetAttrString(cls, "__torch_dispatch__");
  TORCH_CHECK_TYPE(
      attr.ptr() != nullptr &&
          attr.ptr() != torch::disabled_torch_dispatch_impl(),
      ((PyTypeObject*)cls)->tp_name,
      " must define __torch_dispatch__");
#endif

  const auto& local_tensor = r.tensor(1);
  const bool requires_grad = r.toBool(3);
  if (local_tensor.requires_grad() && !requires_grad) {
    TORCH_WARN(
        "To construct DTensor from torch.Tensor, it's recommended to use "
        "local_tensor.detach() and make requires_grad consistent.");
  }
  const auto options = TensorOptions()
                           .dtype(local_tensor.dtype())
                           .device(local_tensor.device())
                           .layout(local_tensor.layout());

  DispatchKeySet extra_dispatch_keys;
  const auto tensor_keys = local_tensor.key_set();
  if (tensor_keys.has(c10::DispatchKey::Conjugate)) {
    extra_dispatch_keys = extra_dispatch_keys.add(c10::DispatchKey::Conjugate);
  }
  if (tensor_keys.has(c10::DispatchKey::Negative)) {
    extra_dispatch_keys = extra_dispatch_keys.add(c10::DispatchKey::Negative);
  }

  py::handle spec = py::handle(r.pyobject(2));
  const auto tensor_meta = spec.attr(dtensor_interned_strings.tensor_meta);
  TORCH_CHECK(!tensor_meta.is_none());
  const auto sizes = tensor_meta.attr(dtensor_interned_strings.shape);
  TORCH_CHECK(
      PyTuple_Check(sizes.ptr()), "spec.tensor_meta.shape must be a tuple");
  const auto stride = tensor_meta.attr(dtensor_interned_strings.stride);
  TORCH_CHECK(
      PyTuple_Check(stride.ptr()), "spec.tensor_meta.stride must be a tuple");

  Tensor tensor = make_tensor_for_subclass_helper(
      /*sym_sizes=*/tuple_to_symintlist(sizes.ptr()),
      /*sym_strides=*/tuple_to_symintlist(stride.ptr()),
      /*sym_storage_offset=*/local_tensor.sym_storage_offset(),
      options,
      /*storage_size=*/std::nullopt,
      extra_dispatch_keys);
  tensor.set_requires_grad(requires_grad);
  tensor.unsafeGetTensorImpl()->set_python_dispatch(true);
  py::object py_tensor = py::reinterpret_steal<py::object>(
      THPVariable_WrapWithType(std::move(tensor), (PyTypeObject*)cls));
  py_tensor.attr(dtensor_interned_strings._spec) = spec;
  py_tensor.attr(dtensor_interned_strings._local_tensor) = local_tensor;
  return py_tensor.release().ptr();
  END_HANDLE_TH_ERRORS
}

struct NativeRuntimeSchemaInfo {
  py::object static_kwargkey;
  size_t static_argnum;
};

NativeRuntimeSchemaInfo unpack_runtime_schema_info(
    py::handle runtime_schema_info,
    size_t num_args) {
  NativeRuntimeSchemaInfo result;
  if (!runtime_schema_info) {
    result.static_argnum = num_args;
  } else {
    result.static_argnum = py::cast<size_t>(
        runtime_schema_info.attr(dtensor_interned_strings.static_argnum));
    result.static_kwargkey =
        runtime_schema_info.attr(dtensor_interned_strings.static_kwargkey);
    TORCH_CHECK(
        result.static_kwargkey.is_none() ||
            PyList_Check(result.static_kwargkey.ptr()),
        "RuntimeSchemaInfo.static_kwargkey must be a list!");
  }
  return result;
}

static bool DTensor_OpSchema_recompute_comparison_key_impl(
    PyObject* self,
    const py::tuple& args_schema) {
  const py::handle self_handle = py::handle(self);
  const auto schema_info =
      self_handle.attr(dtensor_interned_strings.schema_info);
  NativeRuntimeSchemaInfo native_info = unpack_runtime_schema_info(
      checked_not(schema_info.ptr()) ? py::handle() : py::handle(schema_info),
      args_schema.size());
  c10::SmallVector<py::object, 8> args_to_hash;
  size_t idx = 0;
  for (const auto& e : args_schema) {
    if (idx >= native_info.static_argnum ||
        arg_type_tensor_or_tensor_list_like(e)) {
      if (PyList_Check(e.ptr())) {
        args_to_hash.push_back(
            py::reinterpret_steal<py::object>(PyList_AsTuple(e.ptr())));
      } else {
        args_to_hash.push_back(py::reinterpret_borrow<py::object>(e));
      }
    }
    idx++;
  }
  py::tuple args_to_hash_tup(args_to_hash.size());
  for (const auto idx : c10::irange(args_to_hash.size())) {
    args_to_hash_tup[idx] = std::move(args_to_hash[idx]);
  }
  PyObject* comparison_key = nullptr;
  if (native_info.static_kwargkey && !native_info.static_kwargkey.is_none()) {
    py::list static_kwargkey =
        py::reinterpret_borrow<py::list>(native_info.static_kwargkey);
    auto raw_kwargs_schema =
        self_handle.attr(dtensor_interned_strings.kwargs_schema);
    if (!PyDict_Check(raw_kwargs_schema.ptr())) {
      PyErr_SetString(PyExc_TypeError, "self.kwargs_schema must be a dict!");
      return false;
    }
    py::tuple kwargs_to_hash(static_kwargkey.size());
    int idx = 0;
    auto kwargs_schema = py::reinterpret_borrow<py::dict>(raw_kwargs_schema);
    for (const auto& k : static_kwargkey) {
      PyObject* item = PyDict_GetItemWithError(kwargs_schema.ptr(), k.ptr());
      if (item) {
        kwargs_to_hash[idx++] = py::reinterpret_borrow<py::object>(item);
      } else if (PyErr_Occurred()) {
        return false;
      } else {
        kwargs_to_hash[idx++] = py::none();
      }
    }
    comparison_key = PyTuple_Pack(
        3,
        self_handle.attr(dtensor_interned_strings.op).ptr(),
        args_to_hash_tup.ptr(),
        kwargs_to_hash.ptr());
  } else {
    comparison_key = PyTuple_Pack(
        2,
        self_handle.attr(dtensor_interned_strings.op).ptr(),
        args_to_hash_tup.release().ptr());
  }
  if (!comparison_key) {
    return false;
  }
  self_handle.attr(dtensor_interned_strings._comparison_key) =
      py::reinterpret_steal<py::object>(comparison_key);

  return true;
}

static PyObject* DTensor_OpSchema_recompute_comparison_key(
    PyObject* mod,
    PyObject* self) {
  HANDLE_TH_ERRORS
  const py::handle self_handle = py::handle(self);
  const py::handle raw_args_schema =
      self_handle.attr(dtensor_interned_strings.args_schema);
  if (!PyTuple_Check(raw_args_schema.ptr())) {
    PyErr_SetString(PyExc_TypeError, "DTensor.args_schema must be a tuple!");
    return nullptr;
  }
  py::tuple args_schema = py::reinterpret_borrow<py::tuple>(raw_args_schema);
  if (!DTensor_OpSchema_recompute_comparison_key_impl(self, args_schema)) {
    return nullptr;
  }
  Py_RETURN_NONE;
  END_HANDLE_TH_ERRORS
}

static PyObject* DTensor_OpSchema_post_init(PyObject* mod, PyObject* self) {
  HANDLE_TH_ERRORS
  const py::handle self_handle = py::handle(self);
  const py::handle raw_args_schema =
      self_handle.attr(dtensor_interned_strings.args_schema);
  if (!PyTuple_Check(raw_args_schema.ptr())) {
    PyErr_SetString(
        PyExc_TypeError,
        "DTensor_OpSchema_post_init requires self.args_schema to be a tuple!");
    return nullptr;
  }
  py::tuple args_schema = py::reinterpret_borrow<py::tuple>(raw_args_schema);
  if (!DTensor_OpSchema_recompute_comparison_key_impl(self, args_schema)) {
    return nullptr;
  }

  Py_RETURN_NONE;
  END_HANDLE_TH_ERRORS
}

static py::list symint_array_to_list(SymIntArrayRef arr) {
  py::list result(arr.size());
  for (const auto idx : c10::irange(arr.size())) {
    result[idx] = py::cast(arr[idx]);
  }
  return result;
}

static PyObject* DTensor_compute_global_tensor_info_impl(
    const Tensor& tensor,
    py::handle mesh,
    const py::sequence& placements) {
  Py_ssize_t idx = 0;
  c10::SymDimVector tensor_shape(
      tensor.sym_sizes().begin(), tensor.sym_sizes().end());
  c10::SymDimVector tensor_strides(
      tensor.sym_strides().begin(), tensor.sym_strides().end());
  // NOTE: if this is a py::handle then this code stops working;
  // apparently we can't rely on the bound method to stick around.
  py::object mesh_size;
  for (const auto& placement : placements) {
    // TODO: C++ify DeviceMesh somehow; profiling seems
    // to say that nearly all our remaining time spent is spent
    // calling back into Python.
    const auto& cpp_placement = placement.cast<const distributed::Placement&>();
    if (typeid(cpp_placement) == typeid(distributed::Shard) ||
        typeid(cpp_placement) == typeid(distributed::StridedShard)) {
      const auto shard_dim = py::cast<int64_t>(placement.attr("dim"));
      TORCH_CHECK(
          shard_dim >= 0,
          "Shard placements should have negative dims normalized in the user-facing APIs: ",
          py::cast<std::string>(py::str(placement)));
      const auto tensor_ndim = tensor.dim();
      TORCH_CHECK(
          shard_dim < tensor_ndim,
          "Sharding dim ",
          shard_dim,
          " greater than tensor ndim ",
          tensor_ndim,
          " for placement number ",
          idx);

      if (!mesh_size) {
        mesh_size = mesh.attr(dtensor_interned_strings.size);
      }
      const auto mesh_dim_size = py::cast<int64_t>(mesh_size(idx));
      tensor_shape[shard_dim] *= mesh_dim_size;
      // recover tensor stride by modifying the strides that are
      // larger than the current stride on the shard_dim.
      for (const auto i : c10::irange(tensor_strides.size())) {
        if (static_cast<int64_t>(i) != shard_dim &&
            tensor_strides[i] >= tensor_strides[shard_dim]) {
          tensor_strides[i] *= mesh_dim_size;
        }
      }
    } else if (!cpp_placement.is_replicate() && !cpp_placement.is_partial()) {
#if IS_PYTHON_3_11_PLUS
      const auto placement_type_name =
          py::str(py::handle(PyType_GetName(Py_TYPE(placement.ptr()))));
#else
      const auto placement_type_name =
          py::str(py::handle((PyObject*)Py_TYPE(placement.ptr()))
                      .attr(dtensor_interned_strings.__name__));
#endif
      return PyErr_Format(
          PyExc_RuntimeError,
          "placement type %s not supported!",
          py::cast<std::string>(placement_type_name).c_str());
    }
    idx++;
  }
  return py::make_tuple(
             symint_array_to_list(tensor_shape),
             symint_array_to_list(tensor_strides))
      .release()
      .ptr();
}

// NOLINTNEXTLINE(modernize-avoid-c-arrays,cppcoreguidelines-avoid-c-arrays)
static constexpr const char compute_global_tensor_info_doc[] =
    "Compute the global size and stride of a DTensor from the given local tensor.\n"
    "The local size is multiplied by `world_size` per Sharding dim.\n"
    "The local stride is multiplied by `world_size` per Sharding dim, as long as the\n"
    "dimension is outside sharding dim.\n"
    "\n"
    "For example, if we have a local tensor with size (4, 8, 2) and stride (16, 1, 8).\n"
    "If the DTensor placements are [Shard(2)] and world_size is 2;\n"
    "then the global size is (4, 8, 4) and stride is (16 * 2, 1, 8).\n"
    "\n"
    "Args:\n"
    "    tensor (:class:`torch.Tensor`):\n"
    "        Local tensor which DTensor will be constructed from.\n"
    "    mesh (:class:`DeviceMesh`):\n"
    "        Object which describes the mesh topology\n"
    "        of devices for the DTensor.\n"
    "    placements (Sequence[:class:`Placement`]]):\n"
    "        The attribute of the DTensor that describes its layout\n"
    "        on the mesh topology.\n"
    "\n"
    "Return:\n"
    "    tensor_shape: A List of int which specifies the size of DTensor which build\n"
    "        on top of the local tensor.\n"
    "    tensor_stride: A List of int which specifies the stride of DTensor.\n";

static PyObject* DTensor_compute_global_tensor_info(
    PyObject* self,
    PyObject* const* args,
    Py_ssize_t nargs) {
  HANDLE_TH_ERRORS
  TORCH_CHECK_VALUE(
      nargs == 3,
      "compute_global_tensor_info expects 3 arguments, got ",
      nargs);
  TORCH_CHECK_TYPE(
      THPVariable_Check(args[0]),
      "compute_global_tensor_info 1st argument must be Tensor!");
  const auto& tensor = THPVariable_Unpack(args[0]);
  const py::handle mesh = args[1];
  TORCH_CHECK_TYPE(
      PySequence_Check(args[2]),
      "compute_global_tensor_info 3rd argument must be sequence!");
  const py::sequence placements = py::reinterpret_borrow<py::sequence>(args[2]);
  return DTensor_compute_global_tensor_info_impl(tensor, mesh, placements);
  END_HANDLE_TH_ERRORS
}

enum class TensorFlavor {
  NON_TENSOR,
  EXACTLY_DTENSOR,
  EXACTLY_TENSOR,
  DTENSOR_SUBCLASS,
  NON_DTENSOR_TENSOR_SUBCLASS,
};

static std::pair<TensorFlavor, py::object> check_for_dtensor_or_tensor(
    const at::Tensor& tensor) {
  if (!tensor.defined()) {
    return {TensorFlavor::NON_TENSOR, py::object()};
  }

  // I don't think we need to check for wrapped_number() tensors here;
  // the try_replicate_spec_for_scalar_tensor stuff in our caller
  // specifically handles 1-element tensors.

  torch::jit::guardAgainstNamedTensor<at::Tensor>(tensor);
  auto py_tensor = py::cast(tensor);

  const auto dtensor = get_dtensor_class();
  auto* const obj_type = Py_TYPE(py_tensor.ptr());
  if (obj_type == (PyTypeObject*)dtensor.ptr()) {
    return {TensorFlavor::EXACTLY_DTENSOR, std::move(py_tensor)};
  }
  // Fast path for plain old Tensors.
  if (THPVariable_CheckTypeExact(obj_type)) {
    return {TensorFlavor::EXACTLY_TENSOR, std::move(py_tensor)};
  }
  if (py::isinstance(py_tensor, dtensor)) {
    return {TensorFlavor::DTENSOR_SUBCLASS, std::move(py_tensor)};
  }
  TORCH_INTERNAL_ASSERT_DEBUG_ONLY(
      THPVariableClass && py::isinstance(py_tensor, THPVariableClass));
  return {TensorFlavor::NON_DTENSOR_TENSOR_SUBCLASS, std::move(py_tensor)};
}

static std::pair<TensorFlavor, py::object> check_for_dtensor_or_tensor(
    const c10::IValue& iv) {
  if (!iv.isTensor()) {
    return {TensorFlavor::NON_TENSOR, py::object()};
  }

  return check_for_dtensor_or_tensor(iv.toTensor());
}

static c10::List<c10::IValue> replace_dtensors_with_local_tensor(
    const c10::List<c10::IValue>& tl) {
  c10::List<c10::IValue> local_list(tl.elementType());
  local_list.reserve(tl.size());
  for (const auto& elt : tl) {
    const auto [tensor_flavor, py_tensor] = check_for_dtensor_or_tensor(elt);
    if (tensor_flavor == TensorFlavor::EXACTLY_DTENSOR ||
        tensor_flavor == TensorFlavor::DTENSOR_SUBCLASS) {
      local_list.push_back(THPVariable_Unpack(
          py_tensor.attr(dtensor_interned_strings._local_tensor).ptr()));
    } else {
      local_list.push_back(elt);
    }
  }
  return local_list;
}

static void replace_dtensors_with_local_tensor(torch::jit::Stack& stack) {
  for (auto& arg : stack) {
    if (arg.isList()) {
      arg = replace_dtensors_with_local_tensor(arg.toList());
      continue;
    }
    const auto [tensor_flavor, py_tensor] = check_for_dtensor_or_tensor(arg);
    if (tensor_flavor == TensorFlavor::EXACTLY_DTENSOR ||
        tensor_flavor == TensorFlavor::DTENSOR_SUBCLASS) {
      arg = THPVariable_Unpack(
          py_tensor.attr(dtensor_interned_strings._local_tensor).ptr());
    }
  }
}

static py::object try_find_mesh_from_args(
    const c10::OperatorHandle& op,
    const OperatorArgsKwargsView& args_kwargs) {
  for (auto argument_it = args_kwargs.args_begin();
       argument_it != args_kwargs.args_end();
       ++argument_it) {
    const auto [tensor_flavor, py_tensor] =
        check_for_dtensor_or_tensor(*argument_it);
    if (tensor_flavor == TensorFlavor::EXACTLY_DTENSOR ||
        tensor_flavor == TensorFlavor::DTENSOR_SUBCLASS) {
      return py::reinterpret_borrow<py::object>(
          py_tensor.attr(dtensor_interned_strings.device_mesh));
    }
  }
  TORCH_CHECK_VALUE(
      false, "Cannot find device mesh from args for op : ", op.operator_name());
}

static /*DTensorSpec*/ py::object try_replicate_spec_for_scalar_tensor(
    bool allow_implicit_replication,
    py::handle op_call,
    py::handle py_tensor,
    py::handle compute_mesh) {
  const Tensor& tensor_arg = THPVariable_Unpack(py_tensor.ptr());
  const bool numel_is_one = tensor_arg.numel() == 1;
  if (numel_is_one && tensor_arg.dim() == 1) {
    TORCH_WARN(
        "Found a non-scalar tensor with numel=1 and ndim!=0, "
        "we are implicitly creating a replicated DTensor for it. "
        "However, please consider changing it to a scalar tensor "
        "or explicitly create a DTensor under distributed environment.");
  }

  TORCH_CHECK(
      numel_is_one || allow_implicit_replication,
      py::str(op_call),
      " got mixed torch.Tensor and DTensor, need to convert all torch.Tensor to DTensor before calling distributed operators!");

  // scalar tensor can be safely treated as replicated.
  const auto num_placements =
      py::cast<Py_ssize_t>(compute_mesh.attr(dtensor_interned_strings.ndim));
  py::tuple placements_tuple(num_placements);
  py::object replicate = get_replicate_class()();
  for (const auto idx : c10::irange(num_placements)) {
    PyTuple_SET_ITEM(
        placements_tuple.ptr(),
        idx,
        py::reinterpret_borrow<py::object>(replicate).release().ptr());
  }

  return checked_vectorcall(
      get_dtensor_spec_class().ptr(),
      compute_mesh.ptr(),
      placements_tuple.ptr(),
      checked_vectorcall(
          get_tensor_meta_class().ptr(),
          py_tensor.attr(dtensor_interned_strings.shape).ptr(),
          py_tensor.attr(dtensor_interned_strings.stride)().ptr(),
          py_tensor.attr(dtensor_interned_strings.dtype).ptr())
          .ptr());
}

// May return unset object, in which case there was no runtime schema
// info.
static py::object get_runtime_schema_info_for_op(py::handle py_op) {
  const auto op_dispatcher = get_dtensor_op_dispatcher();
  const auto sharding_propagator =
      op_dispatcher.attr(dtensor_interned_strings.sharding_propagator);
  const py::dict op_to_schema_info = py::reinterpret_borrow<py::dict>(
      sharding_propagator.attr(dtensor_interned_strings.op_to_schema_info));
  const py::dict op_to_schema_info_for_single_dim_strategy =
      py::reinterpret_borrow<py::dict>(sharding_propagator.attr(
          dtensor_interned_strings.op_to_schema_info_for_single_dim_strategy));

  PyObject* runtime_schema_info =
      PyDict_GetItemWithError(op_to_schema_info.ptr(), py_op.ptr());
  if (!runtime_schema_info && PyErr_Occurred()) {
    throw py::error_already_set();
  }
  if (!runtime_schema_info) {
    runtime_schema_info = PyDict_GetItemWithError(
        op_to_schema_info_for_single_dim_strategy.ptr(), py_op.ptr());
    if (!runtime_schema_info && PyErr_Occurred()) {
      throw py::error_already_set();
    }
  }
  return py::reinterpret_borrow<py::object>(runtime_schema_info);
}

static bool contains_any_symint(const py::tuple& tup) {
  for (const auto& s : tup) {
    if (THPUtils_checkLong(s.ptr())) {
      continue;
    }
    if (torch::is_symint(s)) {
      return true;
    }
  }
  return false;
}

static bool dtensor_spec_has_symints(py::handle spec) {
  const auto tensor_meta = spec.attr(dtensor_interned_strings.tensor_meta);
  if (tensor_meta.is_none()) {
    return false;
  }
  py::object raw_shape = tensor_meta.attr(dtensor_interned_strings.shape);
  if (!PyTuple_Check(raw_shape.ptr())) {
    PyErr_SetString(PyExc_TypeError, "TensorMeta.shape must be a tuple!");
    throw py::error_already_set();
  }
  const auto shape = py::reinterpret_steal<py::tuple>(raw_shape.release());
  return contains_any_symint(shape);
}

static std::optional<std::pair<NativeOpSchema, /*ComputeMesh*/ py::object>>
create_native_op_schema(
    const c10::OperatorHandle& op,
    py::handle py_op,
    torch::jit::Stack* stack) {
  // fused schema part of unwrap_to_op_info + recompute_comparison_key,
  // operating on IValues instead of Python stuff.

  py::object runtime_schema_info = get_runtime_schema_info_for_op(py_op);
  if (runtime_schema_info &&
      checked_istrue(py::handle(runtime_schema_info)
                         .attr(dtensor_interned_strings.needs_pytree)
                         .ptr())) {
    // Punting on pytree flattening in the fast path on IValues for
    // now since only a minority of ops need it.
    return std::nullopt;
  }

  OperatorArgsKwargsView args_kwargs(op, *stack);
  auto native_info = unpack_runtime_schema_info(
      py::handle(runtime_schema_info), args_kwargs.num_positional_args());

  c10::SmallVector<IValueOrDTensorSpec, 8> comparison_key;
  std::size_t comparison_key_hash = 0;

  py::object compute_mesh = py::none();

  const auto handle_non_dtensor_arg =
      [&comparison_key, &comparison_key_hash, &native_info](
          size_t idx, c10::IValue arg) {
        if (idx >= native_info.static_argnum) {
          if (arg.isList()) {
            const auto& list = arg.toList();
            if (list.empty()) {
              arg = c10::ivalue::Tuple::create({});
            } else {
              // WARNING: here we rely on c10::List being represented
              // by a contiguous array of IValue for efficiency!
              arg = c10::ivalue::Tuple::create(c10::ArrayRef<c10::IValue>(
                  &(*list.begin()).get(), list.size()));
            }
          } else if (arg.isTensor() && !arg.toTensor().defined()) {
            // Coerce undefined Tensor to None, just as we do when
            // converting IValues to PyObject. Otherwise comparison
            // doesn't work. (undefined Tensors can get here because
            // check_for_dtensor_or_tensor calls them non-Tensors, but
            // doesn't have a way to do the coercion for us.)
            arg = c10::IValue();
          }
          comparison_key_hash =
              c10::hash_combine(comparison_key_hash, c10::IValue::hash(arg));
          comparison_key.emplace_back(std::move(arg));
        }
      };
  const auto handle_dtensor_arg = [&comparison_key,
                                   &comparison_key_hash](py::object arg) {
    comparison_key_hash = c10::hash_combine(
        comparison_key_hash, static_cast<size_t>(py::hash(arg)));
    comparison_key.emplace_back(std::move(arg));
  };

  Py_ssize_t idx = 0;
  const bool allow_implicit_replication =
      at::get_dtensor_allow_implicit_replication();
  for (auto argument_it = args_kwargs.args_begin();
       argument_it != args_kwargs.args_end();
       ++argument_it) {
    const auto& arg = *argument_it;
    const auto [tensor_flavor, py_tensor] = check_for_dtensor_or_tensor(arg);
    switch (tensor_flavor) {
      case TensorFlavor::EXACTLY_DTENSOR:
      case TensorFlavor::DTENSOR_SUBCLASS: {
        py::object spec = py_tensor.attr(dtensor_interned_strings._spec);
        if (dtensor_spec_has_symints(spec)) {
          // Symints are unhashable, so we can't use the cache for
          // sharding propagation. bail out to slow path.
          return std::nullopt;
        }
        handle_dtensor_arg(std::move(spec));
        if (compute_mesh.is_none()) {
          compute_mesh = py::reinterpret_borrow<py::object>(
              py_tensor.attr(dtensor_interned_strings.device_mesh));
        }
        break;
      }
      case TensorFlavor::EXACTLY_TENSOR:
      case TensorFlavor::NON_DTENSOR_TENSOR_SUBCLASS: {
        if (compute_mesh.is_none()) {
          compute_mesh = try_find_mesh_from_args(op, args_kwargs);
        }
        handle_dtensor_arg(try_replicate_spec_for_scalar_tensor(
            allow_implicit_replication, py_op, py_tensor, compute_mesh));
        break;
      }
      case TensorFlavor::NON_TENSOR: {
        // Check if this is a list/tuple that might contain DTensors (e.g.,
        // torch.cat)
        if (arg.isList() && compute_mesh.is_none()) {
          const auto list = arg.toList();
          for (const auto& item : list) {
            const auto [item_flavor, item_py_tensor] =
                check_for_dtensor_or_tensor(item);
            if (item_flavor == TensorFlavor::EXACTLY_DTENSOR ||
                item_flavor == TensorFlavor::DTENSOR_SUBCLASS) {
              compute_mesh = py::reinterpret_borrow<py::object>(
                  item_py_tensor.attr(dtensor_interned_strings.device_mesh));
              break;
            }
          }
        }
        // non DTensor/Tensor args (i.e. int/float/bool), just add to
        // local_args
        handle_non_dtensor_arg(idx, arg);
        break;
      }
      default:
        TORCH_INTERNAL_ASSERT(false, "can't happen");
        break;
    }
    idx++;
  }

  // Check kwargs for device_mesh even if not used for cache key
  if (compute_mesh.is_none()) {
    for (auto argument_it = args_kwargs.kwargs_begin();
         argument_it != args_kwargs.kwargs_end();
         ++argument_it) {
      const auto [tensor_flavor, py_tensor] =
          check_for_dtensor_or_tensor(*argument_it);
      if (tensor_flavor == TensorFlavor::EXACTLY_DTENSOR ||
          tensor_flavor == TensorFlavor::DTENSOR_SUBCLASS) {
        compute_mesh = py::reinterpret_borrow<py::object>(
            py_tensor.attr(dtensor_interned_strings.device_mesh));
        break;
      }
    }
  }

  if (native_info.static_kwargkey && !native_info.static_kwargkey.is_none()) {
    // Separator to disambiguate kwargs from args in comparison and hashing.
    static constexpr int64_t kwargs_separator = 0x0011223344556677LL;
    comparison_key.emplace_back(static_cast<int64_t>(kwargs_separator));
    comparison_key_hash = hash_combine(comparison_key_hash, kwargs_separator);

    for (auto argument_it = args_kwargs.kwargs_begin();
         argument_it != args_kwargs.kwargs_end();
         ++argument_it) {
      // Rather than hash/compare the string key, we can just use the
      // index of the kwarg in the schema!
      const auto underlying_index = argument_it.underlying_index();
      comparison_key.emplace_back(c10::IValue(underlying_index));
      comparison_key_hash = hash_combine(
          comparison_key_hash, c10::IValue::hash(comparison_key.back().iv));
      const auto [tensor_flavor, py_tensor] =
          check_for_dtensor_or_tensor(*argument_it);
      switch (tensor_flavor) {
        case TensorFlavor::EXACTLY_DTENSOR:
        case TensorFlavor::DTENSOR_SUBCLASS: {
          handle_dtensor_arg(py_tensor.attr(dtensor_interned_strings._spec));
          if (compute_mesh.is_none()) {
            compute_mesh = py::reinterpret_borrow<py::object>(
                py_tensor.attr(dtensor_interned_strings.device_mesh));
          }
          break;
        }
        case TensorFlavor::EXACTLY_TENSOR:
        case TensorFlavor::NON_DTENSOR_TENSOR_SUBCLASS: {
          handle_dtensor_arg(try_replicate_spec_for_scalar_tensor(
              allow_implicit_replication, py_op, py_tensor, compute_mesh));
          break;
        }
        case TensorFlavor::NON_TENSOR: {
          // Check if this is a list/tuple that might contain DTensors (e.g.,
          // torch.cat)
          if (argument_it->isList() && compute_mesh.is_none()) {
            const auto list = argument_it->toList();
            for (const auto& item : list) {
              const auto [item_flavor, item_py_tensor] =
                  check_for_dtensor_or_tensor(item);
              if (item_flavor == TensorFlavor::EXACTLY_DTENSOR ||
                  item_flavor == TensorFlavor::DTENSOR_SUBCLASS) {
                compute_mesh = py::reinterpret_borrow<py::object>(
                    item_py_tensor.attr(dtensor_interned_strings.device_mesh));
                break;
              }
            }
          }
          handle_non_dtensor_arg(native_info.static_argnum, *argument_it);
          break;
        }
        default:
          TORCH_INTERNAL_ASSERT(false, "can't happen");
          break;
      }
    }
  }

  TORCH_CHECK(
      !compute_mesh.is_none(),
      "found no DeviceMesh from dtensor args for ",
      op.operator_name());

  return std::make_pair(
      NativeOpSchema(
          op,
          std::move(comparison_key),
          comparison_key_hash,
          args_kwargs.num_positional_args()),
      std::move(compute_mesh));
}

static PyObject* get_DTensor_sharding_propagator_cache_stats(
    PyObject* self,
    PyObject* noargs) {
  HANDLE_TH_ERRORS
  auto& cache = get_thread_local_native_sharding_propagator_cache();
  py::tuple result(2);
  result[0] = cache.hits();
  result[1] = cache.misses();
  return result.release().ptr();
  END_HANDLE_TH_ERRORS
}

static PyObject* clear_DTensor_sharding_propagator_cache(
    PyObject* self,
    PyObject* noargs) {
  native_sharding_propagator_cache_DO_NOT_USE.reset();
  Py_RETURN_NONE;
}

using getter = PyObject* (*)(PyObject*, void*);
using setter = int (*)(PyObject*, PyObject*, void*);

static PyObject* THPVariable_get_python_dispatch(
    THPVariable* self,
    void* unused) {
  HANDLE_TH_ERRORS
  const auto& var = THPVariable_Unpack(self);
  return torch::autograd::utils::wrap(
      var.unsafeGetTensorImpl()->is_python_dispatch());
  END_HANDLE_TH_ERRORS
}

// CRTP base class to implement the python bindings for a Tensor property in
// PyTorch A class that implements a property is expected to have:
// - static constexpr const char* name;
//   - This variable should hold the Python name of the property
// - static Tensor fn(const Tensor&);
//   - This function calls the relevant ATen on the tensor
template <typename T>
// NOLINTNEXTLINE(bugprone-crtp-constructor-accessibility)
struct GetterBase {
  static PyObject* getter(THPVariable* self, void* /*unused*/) {
    HANDLE_TH_ERRORS
    if (check_has_torch_function((PyObject*)self)) {
      return handle_torch_function_getter(self, T::name);
    }
    return THPVariable_Wrap(T::fn(THPVariable_Unpack(self)));
    END_HANDLE_TH_ERRORS
  }
};

struct PropertyT : GetterBase<PropertyT> {
  static constexpr const char* name = "T";
  static Tensor fn(const Tensor& t) {
    return t.numpy_T();
  }
};

struct PropertyH : GetterBase<PropertyH> {
  static constexpr const char* name = "H";
  static Tensor fn(const Tensor& t) {
    return t.matrix_H();
  }
};

struct PropertymT : GetterBase<PropertymT> {
  static constexpr const char* name = "mT";
  static Tensor fn(const Tensor& t) {
    return t.mT();
  }
};

struct PropertymH : GetterBase<PropertymH> {
  static constexpr const char* name = "mH";
  static Tensor fn(const Tensor& t) {
    return t.mH();
  }
};

struct PropertyData : GetterBase<PropertyData> {
  static constexpr const char* name = "data";
  static Tensor fn(const Tensor& t) {
    return t.variable_data();
  }
};

struct PropertyGrad : GetterBase<PropertyGrad> {
  static constexpr const char* name = "grad";
  static Tensor fn(const Tensor& t) {
    return t.grad();
  }
};

struct PropertyReal : GetterBase<PropertyReal> {
  static constexpr const char* name = "real";
  static Tensor fn(const Tensor& t) {
    return at::real(t);
  }
};

struct PropertyImag : GetterBase<PropertyImag> {
  static constexpr const char* name = "imag";
  static Tensor fn(const Tensor& t) {
    return at::imag(t);
  }
};

static PyObject* THPVariable_get_cdata(THPVariable* self, void* unused) {
  HANDLE_TH_ERRORS
  if (check_has_torch_function((PyObject*)self)) {
    return handle_torch_function_getter(self, "_cdata");
  }
  const auto& var = THPVariable_Unpack(self);
  return PyLong_FromVoidPtr(var.unsafeGetTensorImpl());
  END_HANDLE_TH_ERRORS
}

static PyObject* THPVariable_get_version(THPVariable* self, void* unused) {
  HANDLE_TH_ERRORS
  if (check_has_torch_function((PyObject*)self)) {
    return handle_torch_function_getter(self, "_version");
  }
  const auto& var = THPVariable_Unpack(self);
  return THPUtils_packInt64(var._version());
  END_HANDLE_TH_ERRORS
}

static PyObject* THPVariable_get_grad_fn(THPVariable* self, void* unused) {
  HANDLE_TH_ERRORS
  if (check_has_torch_function((PyObject*)self)) {
    return handle_torch_function_getter(self, "grad_fn");
  }
  const auto& var = THPVariable_Unpack(self);
  if (!var.grad_fn()) {
    Py_RETURN_NONE;
  }
  return functionToPyObject(var.grad_fn());
  END_HANDLE_TH_ERRORS
}

static int THPVariable_set_grad_fn(
    THPVariable* self,
    PyObject* obj,
    void* unused) {
  HANDLE_TH_ERRORS
  if (check_has_torch_function((PyObject*)self)) {
    return handle_torch_function_setter(self, "_grad_fn", obj);
  }
  TORCH_CHECK(obj, "Deletion of _grad_fn not allowed. Detach tensor instead!");
  TORCH_CHECK(obj == Py_None, "_grad_fn can be only set to None");
  THPVariable_Unpack(self).detach_();
  return 0;
  END_HANDLE_TH_ERRORS_RET(-1)
}

static PyObject* THPVariable_is_leaf(THPVariable* self, void* unused) {
  HANDLE_TH_ERRORS
  if (check_has_torch_function((PyObject*)self)) {
    return handle_torch_function_getter(self, "is_leaf");
  }
  return PyBool_FromLong(!THPVariable_Unpack(self).grad_fn());
  END_HANDLE_TH_ERRORS
}

static int THPVariable_set_data(
    THPVariable* self,
    PyObject* data,
    void* unused) {
  HANDLE_TH_ERRORS
  if (check_has_torch_function((PyObject*)self)) {
    return handle_torch_function_setter(self, "data", data);
  }
  TORCH_CHECK(
      data, "Deleting tensor data is not allowed. Delete tensor instead!");
  TORCH_CHECK_TYPE(
      THPVariable_Check(data),
      "Variable data has to be a tensor, but got ",
      Py_TYPE(data)->tp_name);

  THPVariable_Unpack(self).set_data(THPVariable_Unpack(data));
  return 0;
  END_HANDLE_TH_ERRORS_RET(-1)
}

static int THPVariable_set_grad(
    THPVariable* self,
    PyObject* py_grad,
    void* unused) {
  HANDLE_TH_ERRORS
  if (check_has_torch_function((PyObject*)self)) {
    return handle_torch_function_setter(self, "grad", py_grad);
  }
  const auto& var = THPVariable_Unpack(self);
  if (!py_grad || py_grad == Py_None) {
    var.mutable_grad().reset();
    return 0;
  }

  TORCH_CHECK_TYPE(
      THPVariable_Check(py_grad),
      "assigned grad expected to be a Tensor or None but got grad of type ",
      THPUtils_typename(py_grad));
  TORCH_CHECK(
      self != (THPVariable*)py_grad, "can't assign Variable as its own grad");

  const auto& grad = THPVariable_Unpack(py_grad);
  if (var.grad_dtype().has_value()) {
    TORCH_CHECK(
        grad.dtype() == var.grad_dtype().value(),
        "attempting to assign a gradient with dtype '",
        grad.dtype(),
        "' to a tensor with grad_dtype '",
        var.grad_dtype().value(),
        "'. The gradient must match the tensor's grad_dtype (defaults to the tensor's "
        "dtype). You can set the tensor's grad_dtype attribute with a specific dtype, or "
        "None to allow any dtype. Set grad_dtype with caution. Diverging the dtypes of "
        "a tensor and its gradient may break downstream systems that assume they match.");
  }
  TORCH_CHECK(
      var.device().type() == grad.device().type(),
      "attempting to assign a gradient with device type '",
      grad.device().type(),
      "' to a tensor with device type '",
      var.device().type(),
      "'. Please ensure that the gradient and the tensor are on the same device");
  if (grad.layout() != kSparse) {
    auto expected_options = var.options().dtype(
        var.grad_dtype().has_value() ? var.grad_dtype().value()
                                     : grad.scalar_type());
    TORCH_CHECK(
        grad.options().type_equal(expected_options),
        "attempting to assign a gradient to a tensor that has data of a different type");
  }
  TORCH_CHECK(
      grad.get_device() == var.get_device(),
      "attempting to assign a gradient located on device with index '",
      grad.get_device(),
      "' to a tensor located on device with index '",
      var.get_device(),
      "'. Please ensure that the gradient and the tensor are on the same device");
  TORCH_CHECK(
      grad.sym_sizes().equals(var.sym_sizes()),
      "attempting to assign a gradient of size '",
      grad.sym_sizes(),
      "' to a tensor of size '",
      var.sym_sizes(),
      "'. Please ensure that the gradient and the tensor are the same size");

  var.mutable_grad() = grad;
  return 0;
  END_HANDLE_TH_ERRORS_RET(-1)
}

static PyObject* THPVariable_get_volatile(THPVariable* self, void* unused) {
  HANDLE_TH_ERRORS
  if (check_has_torch_function((PyObject*)self)) {
    return handle_torch_function_getter(self, "volatile");
  }
  const char* msg = "volatile was removed (Variable.volatile is always False)";
  auto r = PyErr_WarnEx(PyExc_UserWarning, msg, 1);
  if (r != 0)
    throw python_error();
  Py_RETURN_FALSE;
  END_HANDLE_TH_ERRORS
}

static int THPVariable_set_volatile(
    THPVariable* self,
    PyObject* obj,
    void* unused) {
  HANDLE_TH_ERRORS
  if (check_has_torch_function((PyObject*)self)) {
    return handle_torch_function_setter(self, "volatile", obj);
  }
  auto r = PyErr_WarnEx(PyExc_UserWarning, VOLATILE_WARNING, 1);
  if (r != 0)
    throw python_error();
  return 0;
  END_HANDLE_TH_ERRORS_RET(-1)
}

static PyObject* THPVariable_get_output_nr(THPVariable* self, void* unused) {
  HANDLE_TH_ERRORS
  if (check_has_torch_function((PyObject*)self)) {
    return handle_torch_function_getter(self, "output_nr");
  }
  const auto output_nr = THPVariable_Unpack(self).output_nr();
  return THPUtils_packInt64(output_nr);
  END_HANDLE_TH_ERRORS
}

static PyObject* THPVariable_get_requires_grad(
    THPVariable* self,
    void* unused) {
  HANDLE_TH_ERRORS
  if (check_has_torch_function((PyObject*)self)) {
    return handle_torch_function_getter(self, "requires_grad");
  }
  if (THPVariable_Unpack(self).requires_grad()) {
    Py_RETURN_TRUE;
  } else {
    Py_RETURN_FALSE;
  }
  END_HANDLE_TH_ERRORS
}

static PyObject* THPVariable_retains_grad(THPVariable* self, void* unused) {
  HANDLE_TH_ERRORS
  if (check_has_torch_function((PyObject*)self)) {
    return handle_torch_function_getter(self, "retains_grad");
  }
  if (THPVariable_Unpack(self).retains_grad()) {
    Py_RETURN_TRUE;
  } else {
    Py_RETURN_FALSE;
  }
  END_HANDLE_TH_ERRORS
}

static PyObject* THPVariable_get_ndim(THPVariable* self, void* unused) {
  HANDLE_TH_ERRORS
  if (check_has_torch_function((PyObject*)self)) {
    return handle_torch_function_getter(self, "ndim");
  }
  return THPUtils_packInt64(THPVariable_Unpack(self).dim());
  END_HANDLE_TH_ERRORS
}

static PyObject* THPVariable_get_names(PyObject* self, void* unused) {
  HANDLE_TH_ERRORS
  if (check_has_torch_function(self)) {
    return handle_torch_function_getter((THPVariable*)self, "names");
  }
  // The long-term plan is to return a list of (python) torch.Dimname.
  // However, for now, return a list of string.
  const auto& tensor = THPVariable_Unpack(self);
  auto size = tensor.dim();
  THPObjectPtr tuple(PyTuple_New(size));
  if (!tuple)
    throw python_error();

  const auto dimnames = tensor.names();
  for (const auto i : c10::irange(size)) {
    PyObject* str = nullptr;
    if (dimnames[i].type() == at::NameType::WILDCARD) {
      // PyTuple_SET_ITEM steals a reference to the object. When the tuple is
      // deallocated, it'll decrement the refcount on Py_None, which is bad.
      // To avoid this, we "create" a new reference to Py_None by increasing
      // the refcount.
      // Sources:
      // - https://docs.python.org/3/c-api/tuple.html#c.PyTuple_SetItem
      // -
      // https://stackoverflow.com/questions/16400600/how-to-return-a-tuple-containing-a-none-value-from-the-c-api
      Py_INCREF(Py_None);
      str = Py_None;
    } else {
      str = THPUtils_packString(dimnames[i].symbol().toUnqualString());
      if (!str)
        throw python_error();
    }
    PyTuple_SET_ITEM(tuple.get(), i, str);
  }
  return tuple.release();
  END_HANDLE_TH_ERRORS
}

static int THPVariable_set_names(
    PyObject* self,
    PyObject* names,
    void* unused) {
  HANDLE_TH_ERRORS
  if (check_has_torch_function(self)) {
    return handle_torch_function_setter((THPVariable*)self, "names", names);
  }
  const auto& var = THPVariable_Unpack(self);
  if (names == Py_None) {
    at::internal_set_names_inplace(var, std::nullopt);
  } else {
    TORCH_CHECK(
        THPUtils_checkDimnameList(names),
        "names must either be None or a tuple of dim names");
    at::internal_set_names_inplace(var, torch::parseDimnameList(names));
  }
  return 0;
  END_HANDLE_TH_ERRORS_RET(-1)
}

static int THPVariable_set_requires_grad(
    THPVariable* self,
    PyObject* obj,
    void* unused) {
  HANDLE_TH_ERRORS
  if (check_has_torch_function((PyObject*)self)) {
    return handle_torch_function_setter(self, "requires_grad", obj);
  }
  TORCH_CHECK(obj && PyBool_Check(obj), "requires_grad must be a bool");
  const auto& var = THPVariable_Unpack(self);
  auto requires_grad = (obj == Py_True);
  if (!var.is_leaf()) {
    THPUtils_setError(
        autograd::utils::requires_grad_leaf_error(obj == Py_True).c_str());
    return -1;
  }
  if (requires_grad &&
      !isDifferentiableType(at::typeMetaToScalarType((var.dtype())))) {
    THPUtils_setError(
        "only Tensors of floating point and complex dtype can require gradients");
    return -1;
  }
  var.set_requires_grad(requires_grad);
  return 0;
  END_HANDLE_TH_ERRORS_RET(-1)
}

static PyObject* THPVariable_get_name(THPVariable* self, void* unused) {
  if (check_has_torch_function((PyObject*)self)) {
    HANDLE_TH_ERRORS
    return handle_torch_function_getter(self, "name");
    END_HANDLE_TH_ERRORS
  }
  const auto& tensor = THPVariable_Unpack(self);
  if (tensor.name().empty())
    Py_RETURN_NONE;
  return THPUtils_packString(tensor.name().c_str());
}

static PyObject* THPVariable_get_backwards_hooks(
    THPVariable* self,
    void* unused) {
  HANDLE_TH_ERRORS
  if (check_has_torch_function((PyObject*)self)) {
    return handle_torch_function_getter(self, "_backward_hooks");
  }
  if (self->backward_hooks) {
    Py_INCREF(self->backward_hooks);
    return self->backward_hooks;
  }
  Py_RETURN_NONE;
  END_HANDLE_TH_ERRORS
}

static int THPVariable_set_backwards_hooks(
    THPVariable* self,
    PyObject* obj,
    void* unused) {
  HANDLE_TH_ERRORS
  if (check_has_torch_function((PyObject*)self)) {
    return handle_torch_function_setter(self, "_backward_hooks", obj);
  }
  TORCH_CHECK(obj, "Deletion of _backwards_hooks not allowed!");
  if (obj == Py_None) {
    obj = nullptr;
  }
  Py_XINCREF(obj);
  Py_XDECREF(self->backward_hooks);
  self->backward_hooks = obj;
  const auto& tensor = THPVariable_Unpack(self);
  torch::autograd::impl::clear_hooks(tensor);
  if (obj) {
    torch::autograd::impl::add_hook(
        tensor, std::make_unique<PyFunctionTensorPreHook>(obj, 0));
  }
  return 0;
  END_HANDLE_TH_ERRORS_RET(-1)
}

static PyObject* THPVariable_get_post_accumulate_grad_hooks(
    THPVariable* self,
    void* unused) {
  HANDLE_TH_ERRORS
  if (check_has_torch_function((PyObject*)self)) {
    return handle_torch_function_getter(self, "_post_accumulate_grad_hooks");
  }
  if (self->post_accumulate_grad_hooks) {
    Py_INCREF(self->post_accumulate_grad_hooks);
    return self->post_accumulate_grad_hooks;
  }
  Py_RETURN_NONE;
  END_HANDLE_TH_ERRORS
}

static int THPVariable_set_post_accumulate_grad_hooks(
    THPVariable* self,
    PyObject* obj,
    void* unused) {
  HANDLE_TH_ERRORS
  if (check_has_torch_function((PyObject*)self)) {
    return handle_torch_function_setter(
        self, "_post_accumulate_grad_hooks", obj);
  }
  TORCH_CHECK(obj, "Deletion of _post_accumulate_grad_hooks not allowed!");
  if (obj == Py_None) {
    obj = nullptr;
  }
  Py_XINCREF(obj);
  Py_CLEAR(self->post_accumulate_grad_hooks);
  self->post_accumulate_grad_hooks = obj;
  const auto& tensor = THPVariable_Unpack(self);
  if (obj) {
    torch::autograd::impl::set_post_acc_grad_hooks(
        tensor, std::make_unique<PyFunctionTensorPostAccGradHooks>(obj));
  }
  return 0;
  END_HANDLE_TH_ERRORS_RET(-1)
}

static PyObject* THPVariable_get_base(THPVariable* self, void* unused) {
  HANDLE_TH_ERRORS
  if (check_has_torch_function((PyObject*)self)) {
    return handle_torch_function_getter(self, "_base");
  }
  const auto& tensor = THPVariable_Unpack(self);
  if (tensor.is_view()) {
    return THPVariable_Wrap(tensor._base());
  }
  Py_RETURN_NONE;
  END_HANDLE_TH_ERRORS
}

static PyObject* THPVariable_get_shape(THPVariable* self, void* unused) {
  HANDLE_TH_ERRORS
  if (check_has_torch_function((PyObject*)self)) {
    return handle_torch_function_getter(self, "shape");
  }
  return THPSize_NewFromSymSizes(THPVariable_Unpack(self));
  END_HANDLE_TH_ERRORS
}

static PyObject* THPVariable_is_cpu(THPVariable* self, void* unused) {
  HANDLE_TH_ERRORS
  if (check_has_torch_function((PyObject*)self)) {
    return handle_torch_function_getter(self, "is_cpu");
  }
  auto& self_ = THPVariable_Unpack(self);
  return torch::autograd::utils::wrap(self_.is_cpu());
  END_HANDLE_TH_ERRORS
}

static PyObject* THPVariable_is_cuda(THPVariable* self, void* unused) {
  HANDLE_TH_ERRORS
  if (check_has_torch_function((PyObject*)self)) {
    return handle_torch_function_getter(self, "is_cuda");
  }
  auto& self_ = THPVariable_Unpack(self);
  return torch::autograd::utils::wrap(self_.is_cuda());
  END_HANDLE_TH_ERRORS
}

static PyObject* THPVariable_is_mtia(THPVariable* self, void* unused) {
  HANDLE_TH_ERRORS
  if (check_has_torch_function((PyObject*)self)) {
    return handle_torch_function_getter(self, "is_mtia");
  }
  auto& self_ = THPVariable_Unpack(self);
  return torch::autograd::utils::wrap(self_.is_mtia());
  END_HANDLE_TH_ERRORS
}

static PyObject* THPVariable_is_xla(THPVariable* self, void* unused) {
  HANDLE_TH_ERRORS
  if (check_has_torch_function((PyObject*)self)) {
    return handle_torch_function_getter(self, "is_xla");
  }
  auto& self_ = THPVariable_Unpack(self);
  return torch::autograd::utils::wrap(self_.is_xla());
  END_HANDLE_TH_ERRORS
}

static PyObject* THPVariable_is_ipu(THPVariable* self, void* unused) {
  HANDLE_TH_ERRORS
  if (check_has_torch_function((PyObject*)self)) {
    return handle_torch_function_getter(self, "is_ipu");
  }
  auto& self_ = THPVariable_Unpack(self);
  return torch::autograd::utils::wrap(self_.is_ipu());
  END_HANDLE_TH_ERRORS
}

static PyObject* THPVariable_is_xpu(THPVariable* self, void* unused) {
  HANDLE_TH_ERRORS
  if (check_has_torch_function((PyObject*)self)) {
    return handle_torch_function_getter(self, "is_xpu");
  }
  auto& self_ = THPVariable_Unpack(self);
  return torch::autograd::utils::wrap(self_.is_xpu());
  END_HANDLE_TH_ERRORS
}

static PyObject* THPVariable_is_sparse(THPVariable* self, void* unused) {
  HANDLE_TH_ERRORS
  if (check_has_torch_function((PyObject*)self)) {
    return handle_torch_function_getter(self, "is_sparse");
  }
  auto& self_ = THPVariable_Unpack(self);
  return torch::autograd::utils::wrap(self_.is_sparse());
  END_HANDLE_TH_ERRORS
}

static PyObject* THPVariable_is_sparse_csr(THPVariable* self, void* unused) {
  HANDLE_TH_ERRORS
  if (check_has_torch_function((PyObject*)self)) {
    return handle_torch_function_getter(self, "is_sparse_csr");
  }
  auto& self_ = THPVariable_Unpack(self);
  return torch::autograd::utils::wrap(self_.is_sparse_csr());
  END_HANDLE_TH_ERRORS
}

static PyObject* THPVariable_is_mkldnn(THPVariable* self, void* unused) {
  HANDLE_TH_ERRORS
  if (check_has_torch_function((PyObject*)self)) {
    return handle_torch_function_getter(self, "is_mkldnn");
  }
  auto& self_ = THPVariable_Unpack(self);
  return torch::autograd::utils::wrap(self_.is_mkldnn());
  END_HANDLE_TH_ERRORS
}

static PyObject* THPVariable_is_mps(THPVariable* self, void* unused) {
  HANDLE_TH_ERRORS
  if (check_has_torch_function((PyObject*)self)) {
    return handle_torch_function_getter(self, "is_mps");
  }
  auto& self_ = THPVariable_Unpack(self);
  return torch::autograd::utils::wrap(self_.is_mps());
  END_HANDLE_TH_ERRORS
}

static PyObject* THPVariable_is_maia(THPVariable* self, void* unused) {
  HANDLE_TH_ERRORS
  if (check_has_torch_function((PyObject*)self)) {
    return handle_torch_function_getter(self, "is_maia");
  }
  auto& self_ = THPVariable_Unpack(self);
  return torch::autograd::utils::wrap(self_.is_maia());
  END_HANDLE_TH_ERRORS
}

static PyObject* THPVariable_is_vulkan(THPVariable* self, void* unused) {
  HANDLE_TH_ERRORS
  if (check_has_torch_function((PyObject*)self)) {
    return handle_torch_function_getter(self, "is_vulkan");
  }
  auto& self_ = THPVariable_Unpack(self);
  return torch::autograd::utils::wrap(self_.is_vulkan());
  END_HANDLE_TH_ERRORS
}

static PyObject* THPVariable_is_quantized(THPVariable* self, void* unused) {
  HANDLE_TH_ERRORS
  if (check_has_torch_function((PyObject*)self)) {
    return handle_torch_function_getter(self, "is_quantized");
  }
  auto& self_ = THPVariable_Unpack(self);
  return torch::autograd::utils::wrap(self_.is_quantized());
  END_HANDLE_TH_ERRORS
}

static PyObject* THPVariable_is_meta(THPVariable* self, void* unused) {
  HANDLE_TH_ERRORS
  if (check_has_torch_function((PyObject*)self)) {
    return handle_torch_function_getter(self, "is_meta");
  }
  auto& self_ = THPVariable_Unpack(self);
  return torch::autograd::utils::wrap(self_.is_meta());
  END_HANDLE_TH_ERRORS
}

static PyObject* THPVariable_is_complex(THPVariable* self, void* unused) {
  HANDLE_TH_ERRORS
  if (check_has_torch_function((PyObject*)self)) {
    return handle_torch_function_getter(self, "is_complex");
  }
  auto& self_ = THPVariable_Unpack(self);
  return torch::autograd::utils::wrap(self_.is_complex());
  END_HANDLE_TH_ERRORS
}

static PyObject* THPVariable_is_nested(THPVariable* self, void* unused) {
  HANDLE_TH_ERRORS
  if (check_has_torch_function((PyObject*)self)) {
    return handle_torch_function_getter(self, "is_nested");
  }
  auto& self_ = THPVariable_Unpack(self);
  return torch::autograd::utils::wrap(self_.is_nested());
  END_HANDLE_TH_ERRORS
}

static PyObject* THPVariable_has_symbolic_sizes_strides(
    THPVariable* self,
    void* unused) {
  HANDLE_TH_ERRORS
  auto& self_ = THPVariable_Unpack(self);
  return torch::autograd::utils::wrap(
      self_.unsafeGetTensorImpl()->has_symbolic_sizes_strides());
  END_HANDLE_TH_ERRORS
}

static PyObject* THPVariable_dtype(THPVariable* self, void* unused) {
  HANDLE_TH_ERRORS
  if (check_has_torch_function((PyObject*)self)) {
    return handle_torch_function_getter(self, "dtype");
  }
  auto& self_ = THPVariable_Unpack(self);
  return torch::autograd::utils::wrap(self_.scalar_type());
  END_HANDLE_TH_ERRORS
}

static PyObject* THPVariable_layout(THPVariable* self, void* unused) {
  HANDLE_TH_ERRORS
  if (check_has_torch_function((PyObject*)self)) {
    return handle_torch_function_getter(self, "layout");
  }
  auto& self_ = THPVariable_Unpack(self);
  return torch::autograd::utils::wrap(self_.layout());
  END_HANDLE_TH_ERRORS
}

static PyObject* THPVariable_device(THPVariable* self, void* unused) {
  HANDLE_TH_ERRORS
  if (check_has_torch_function((PyObject*)self)) {
    return handle_torch_function_getter(self, "device");
  }
  return THPDevice_New(THPVariable_Unpack(self).device());
  END_HANDLE_TH_ERRORS
}

static PyObject* THPVariable_get_nbytes(THPVariable* self, void* unused) {
  HANDLE_TH_ERRORS
  if (check_has_torch_function((PyObject*)self)) {
    return handle_torch_function_getter(self, "nbytes");
  }
  return PyLong_FromSize_t(THPVariable_Unpack(self).nbytes());
  END_HANDLE_TH_ERRORS
}

static PyObject* THPVariable_get_grad_dtype(THPVariable* self, void* unused) {
  HANDLE_TH_ERRORS
  if (check_has_torch_function((PyObject*)self)) {
    return handle_torch_function_getter(self, "grad_dtype");
  }
  const auto& var = THPVariable_Unpack(self);
  TORCH_CHECK(
      !var.grad_fn(), "grad_dtype can only be accessed on leaf tensors.");
  if (!var.grad_dtype().has_value()) {
    Py_RETURN_NONE;
  } else {
    return torch::autograd::utils::wrap(var.grad_dtype().value());
  }
  END_HANDLE_TH_ERRORS
}

static int THPVariable_set_grad_dtype(
    THPVariable* self,
    PyObject* obj,
    void* unused) {
  HANDLE_TH_ERRORS
  if (check_has_torch_function((PyObject*)self)) {
    return handle_torch_function_setter(self, "grad_dtype", obj);
  }
  const auto& var = THPVariable_Unpack(self);
  TORCH_CHECK(
      THPDtype_Check(obj) || obj == Py_None,
      "grad_dtype must be a torch.dtype or None, but got ",
      Py_TYPE(obj)->tp_name);
  if (var.grad().defined() && obj != Py_None) {
    auto new_dtype = reinterpret_cast<THPDtype*>(obj);
    TORCH_CHECK(
        var.grad().dtype() == new_dtype->scalar_type,
        "Cannot set grad_dtype to '",
        new_dtype->scalar_type,
        "' because there is already a gradient with dtype '",
        var.grad().dtype(),
        "'. Please clear the gradient (.grad = None) before changing grad_dtype, "
        "or ensure the new grad_dtype matches the existing gradient's dtype.");
  }
  std::optional<at::ScalarType> new_dtype;
  if (obj != Py_None) {
    auto* dtype = reinterpret_cast<THPDtype*>(obj);
    new_dtype = dtype->scalar_type;
  }
  var.set_grad_dtype(new_dtype);
  return 0;
  END_HANDLE_TH_ERRORS_RET(-1)
}

static PyObject* THPVariable_get_itemsize(THPVariable* self, void* unused) {
  HANDLE_TH_ERRORS
  if (check_has_torch_function((PyObject*)self)) {
    return handle_torch_function_getter(self, "itemsize");
  }
  return PyLong_FromSize_t(THPVariable_Unpack(self).itemsize());
  END_HANDLE_TH_ERRORS
}

static int THPVariable_set_real(PyObject* self, PyObject* real, void* unused) {
  HANDLE_TH_ERRORS
  auto& self_ = THPVariable_Unpack(self);
  auto self_real = at::real(self_);
  auto real_ = valueToTensor(self_real.options(), real, self_real.device());
  {
    pybind11::gil_scoped_release no_gil;
    self_real.copy_(real_);
    return 0;
  }
  END_HANDLE_TH_ERRORS_RET(-1)
}

static int THPVariable_set_imag(PyObject* self, PyObject* imag, void* unused) {
  HANDLE_TH_ERRORS
  auto& self_ = THPVariable_Unpack(self);
  auto self_imag = at::imag(self_);
  auto imag_ = valueToTensor(self_imag.options(), imag, self_imag.device());
  {
    pybind11::gil_scoped_release no_gil;
    self_imag.copy_(imag_);
    return 0;
  }
  END_HANDLE_TH_ERRORS_RET(-1)
}

static PyObject* THPVariable__use_count(PyObject* self, PyObject* noargs) {
  HANDLE_TH_ERRORS
  const auto& t = THPVariable_Unpack(self);
  return THPUtils_packUInt64(t.use_count());
  END_HANDLE_TH_ERRORS
}

// properties are registered here because we are currently only able to bind
// them manually. TODO: make declarable in native_functions
// NOLINTNEXTLINE(modernize-avoid-c-arrays,cppcoreguidelines-avoid-c-arrays,cppcoreguidelines-avoid-non-const-global-variables)
static struct PyGetSetDef THPVariable_properties[] = {
    {"_python_dispatch",
     (getter)THPVariable_get_python_dispatch,
     nullptr,
     nullptr,
     nullptr},
    {"T", (getter)PropertyT::getter, nullptr, nullptr, nullptr},
    {"H", (getter)PropertyH::getter, nullptr, nullptr, nullptr},
    {"mT", (getter)PropertymT::getter, nullptr, nullptr, nullptr},
    {"mH", (getter)PropertymH::getter, nullptr, nullptr, nullptr},
    {"_cdata", (getter)THPVariable_get_cdata, nullptr, nullptr, nullptr},
    {"_version", (getter)THPVariable_get_version, nullptr, nullptr, nullptr},
    {"grad_fn", (getter)THPVariable_get_grad_fn, nullptr, nullptr, nullptr},
    {"_grad_fn",
     (getter)THPVariable_get_grad_fn,
     (setter)THPVariable_set_grad_fn,
     nullptr,
     nullptr},
    {"is_leaf", (getter)THPVariable_is_leaf, nullptr, nullptr, nullptr},
    {"retains_grad",
     (getter)THPVariable_retains_grad,
     nullptr,
     nullptr,
     nullptr},
    {"data",
     (getter)PropertyData::getter,
     (setter)THPVariable_set_data,
     nullptr,
     nullptr},
    {"_grad",
     (getter)PropertyGrad::getter,
     (setter)THPVariable_set_grad,
     nullptr,
     nullptr}, // Allows the python class to override .grad
    {"grad",
     (getter)PropertyGrad::getter,
     (setter)THPVariable_set_grad,
     nullptr,
     nullptr},
    {"_base", (getter)THPVariable_get_base, nullptr, nullptr, nullptr},
    {"volatile",
     (getter)THPVariable_get_volatile,
     (setter)THPVariable_set_volatile,
     nullptr,
     nullptr},
    {"output_nr", (getter)THPVariable_get_output_nr, nullptr, nullptr, nullptr},
    {"requires_grad",
     (getter)THPVariable_get_requires_grad,
     (setter)THPVariable_set_requires_grad,
     nullptr,
     nullptr},
    {"_backward_hooks",
     (getter)THPVariable_get_backwards_hooks,
     (setter)THPVariable_set_backwards_hooks,
     nullptr,
     nullptr},
    {"_post_accumulate_grad_hooks",
     (getter)THPVariable_get_post_accumulate_grad_hooks,
     (setter)THPVariable_set_post_accumulate_grad_hooks,
     nullptr,
     nullptr},
    {"name", (getter)THPVariable_get_name, nullptr, nullptr, nullptr},
    {"shape", (getter)THPVariable_get_shape, nullptr, nullptr, nullptr},
    {"is_cuda", (getter)THPVariable_is_cuda, nullptr, nullptr, nullptr},
    {"is_mtia", (getter)THPVariable_is_mtia, nullptr, nullptr, nullptr},
    {"is_cpu", (getter)THPVariable_is_cpu, nullptr, nullptr, nullptr},
    {"is_xla", (getter)THPVariable_is_xla, nullptr, nullptr, nullptr},
    {"is_xpu", (getter)THPVariable_is_xpu, nullptr, nullptr, nullptr},
    {"is_ipu", (getter)THPVariable_is_ipu, nullptr, nullptr, nullptr},
    {"is_sparse", (getter)THPVariable_is_sparse, nullptr, nullptr, nullptr},
    {"is_sparse_csr",
     (getter)THPVariable_is_sparse_csr,
     nullptr,
     nullptr,
     nullptr},
    {"is_mkldnn", (getter)THPVariable_is_mkldnn, nullptr, nullptr, nullptr},
    {"is_mps", (getter)THPVariable_is_mps, nullptr, nullptr, nullptr},
    {"is_maia", (getter)THPVariable_is_maia, nullptr, nullptr, nullptr},
    {"is_vulkan", (getter)THPVariable_is_vulkan, nullptr, nullptr, nullptr},
    {"is_complex", (getter)THPVariable_is_complex, nullptr, nullptr, nullptr},
    {"is_quantized",
     (getter)THPVariable_is_quantized,
     nullptr,
     nullptr,
     nullptr},
    {"is_meta", (getter)THPVariable_is_meta, nullptr, nullptr, nullptr},
    {"is_nested", (getter)THPVariable_is_nested, nullptr, nullptr, nullptr},
    {"_has_symbolic_sizes_strides",
     (getter)THPVariable_has_symbolic_sizes_strides,
     nullptr,
     nullptr,
     nullptr},
    {"dtype", (getter)THPVariable_dtype, nullptr, nullptr, nullptr},
    {"layout", (getter)THPVariable_layout, nullptr, nullptr, nullptr},
    {"device", (getter)THPVariable_device, nullptr, nullptr, nullptr},
    {"ndim", (getter)THPVariable_get_ndim, nullptr, nullptr, nullptr},
    {"nbytes", (getter)THPVariable_get_nbytes, nullptr, nullptr, nullptr},
    {"itemsize", (getter)THPVariable_get_itemsize, nullptr, nullptr, nullptr},
    {"names",
     (getter)THPVariable_get_names,
     (setter)THPVariable_set_names,
     nullptr,
     nullptr},
    {"real",
     (getter)PropertyReal::getter,
     (setter)THPVariable_set_real,
     nullptr,
     nullptr},
    {"imag",
     (getter)PropertyImag::getter,
     (setter)THPVariable_set_imag,
     nullptr,
     nullptr},
    {"grad_dtype",
     (getter)THPVariable_get_grad_dtype,
     (setter)THPVariable_set_grad_dtype,
     nullptr,
     nullptr},
    {nullptr}};

static PyMappingMethods THPVariable_as_mapping = {
    THPVariable_length,
    THPVariable_getitem,
    THPVariable_setitem,
};

// NOLINTNEXTLINE(modernize-avoid-c-arrays,cppcoreguidelines-avoid-c-arrays,cppcoreguidelines-avoid-non-const-global-variables)
static PyMethodDef extra_methods[] = {
    {"as_subclass",
     castPyCFunctionWithKeywords(THPVariable_as_subclass),
     METH_VARARGS | METH_KEYWORDS,
     nullptr},
    {"_make_subclass",
     castPyCFunctionWithKeywords(THPVariable_make_subclass),
     METH_STATIC | METH_VARARGS | METH_KEYWORDS,
     nullptr},
    {"_make_wrapper_subclass",
     castPyCFunctionWithKeywords(THPVariable_make_wrapper_subclass),
     METH_STATIC | METH_VARARGS | METH_KEYWORDS,
     nullptr},
    {"_dtensor__new__",
     castPyCFunctionWithKeywords(THPVariable_dtensor_new),
     METH_STATIC | METH_VARARGS | METH_KEYWORDS,
     nullptr},
    {"_fix_weakref", THPVariable_fix_weakref, METH_NOARGS, nullptr},
    {"_view_func",
     castPyCFunctionWithKeywords(THPVariable_view_func),
     METH_VARARGS | METH_KEYWORDS,
     nullptr},
    {"_view_func_unsafe",
     castPyCFunctionWithKeywords(THPVariable_view_func_unsafe),
     METH_VARARGS | METH_KEYWORDS,
     nullptr},
    {"_rev_view_func_unsafe",
     THPVariable_rev_view_func_unsafe,
     METH_O,
     nullptr},
    {"_use_count", THPVariable__use_count, METH_NOARGS, nullptr},
    {nullptr}};

// NOLINTNEXTLINE(modernize-avoid-c-arrays,cppcoreguidelines-avoid-c-arrays,cppcoreguidelines-avoid-non-const-global-variables)
static PyMethodDef extra_dtensor_functions[] = {
    {"_DTensor_OpSchema_post_init",
     DTensor_OpSchema_post_init,
     METH_O,
     nullptr},
    {"_DTensor_OpSchema_recompute_comparison_key",
     DTensor_OpSchema_recompute_comparison_key,
     METH_O,
     nullptr},
    {"_DTensor_compute_global_tensor_info",
     castPyCFunctionFast(DTensor_compute_global_tensor_info),
     METH_FASTCALL,
     compute_global_tensor_info_doc},
    {"_get_DTensor_sharding_propagator_cache_stats",
     get_DTensor_sharding_propagator_cache_stats,
     METH_NOARGS,
     nullptr},
    {"_clear_DTensor_sharding_propagator_cache",
     clear_DTensor_sharding_propagator_cache,
     METH_NOARGS,
     nullptr},
    {nullptr}};

struct THPVariableMeta {
  PyHeapTypeObject base;
};

static int THPVariableMetaType_init(
    PyObject* cls,
    PyObject* args,
    PyObject* kwargs);

static PyTypeObject THPVariableMetaType = {
    PyVarObject_HEAD_INIT(DEFERRED_ADDRESS(&PyType_Type), 0)
    "torch._C._TensorMeta", /* tp_name */
    sizeof(THPVariableMeta), /* tp_basicsize */
    0, /* tp_itemsize */
    nullptr, /* tp_dealloc */
    0, /* tp_vectorcall_offset */
    nullptr, /* tp_getattr */
    nullptr, /* tp_setattr */
    nullptr, /* tp_reserved */
    nullptr, /* tp_repr */
    nullptr, /* tp_as_number */
    nullptr, /* tp_as_sequence */
    nullptr, /* tp_as_mapping */
    nullptr, /* tp_hash  */
    nullptr, /* tp_call */
    nullptr, /* tp_str */
    nullptr, /* tp_getattro */
    nullptr, /* tp_setattro */
    nullptr, /* tp_as_buffer */
    // NOLINTNEXTLINE(misc-redundant-expression)
    Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE, /* tp_flags */
    nullptr, /* tp_doc */
    nullptr, /* tp_traverse */
    nullptr, /* tp_clear */
    nullptr, /* tp_richcompare */
    0, /* tp_weaklistoffset */
    nullptr, /* tp_iter */
    nullptr, /* tp_iternext */
    nullptr, /* tp_methods */
    nullptr, /* tp_members */
    nullptr, /* tp_getset */
    DEFERRED_ADDRESS(&PyType_Type), /* tp_base */
    nullptr, /* tp_dict */
    nullptr, /* tp_descr_get */
    nullptr, /* tp_descr_set */
    0, /* tp_dictoffset */
    THPVariableMetaType_init, /* tp_init */
    nullptr, /* tp_alloc */
    nullptr, /* tp_new */
};

static void THPVariable_dealloc(PyObject* self);
static int THPVariable_clear(THPVariable* self);
static int THPVariable_traverse(PyObject* self, visitproc visit, void* arg);

static PyTypeObject THPVariableType = {
    PyVarObject_HEAD_INIT(&THPVariableMetaType, 0)
    "torch._C.TensorBase", /* tp_name */
    sizeof(THPVariable), /* tp_basicsize */
    0, /* tp_itemsize */
    THPVariable_dealloc, /* tp_dealloc */
    0, /* tp_vectorcall_offset */
    nullptr, /* tp_getattr */
    nullptr, /* tp_setattr */
    nullptr, /* tp_reserved */
    nullptr, /* tp_repr */
    nullptr, /* tp_as_number */
    nullptr, /* tp_as_sequence */
    &THPVariable_as_mapping, /* tp_as_mapping */
    nullptr, /* tp_hash  */
    nullptr, /* tp_call */
    nullptr, /* tp_str */
    nullptr, /* tp_getattro */
    nullptr, /* tp_setattro */
    nullptr, /* tp_as_buffer */
    // NOLINTNEXTLINE(misc-redundant-expression)
    Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE |
        Py_TPFLAGS_HAVE_GC, /* tp_flags */
    nullptr, /* tp_doc */
    (traverseproc)THPVariable_traverse, /* tp_traverse */
    (inquiry)THPVariable_clear, /* tp_clear */
    nullptr, /* tp_richcompare */
    0, /* tp_weaklistoffset */
    nullptr, /* tp_iter */
    nullptr, /* tp_iternext */
    nullptr, /* tp_methods */
    nullptr, /* tp_members */
    THPVariable_properties, /* tp_getset */
    nullptr, /* tp_base */
    nullptr, /* tp_dict */
    nullptr, /* tp_descr_get */
    nullptr, /* tp_descr_set */
    0, /* tp_dictoffset */
    nullptr, /* tp_init */
    nullptr, /* tp_alloc */
    // Although new is provided here, it is illegal to call this with cls ==
    // THPVariableMeta.  Instead, subclass it first and then construct it
    THPVariable_pynew, /* tp_new */
};

PyObject* THPVariable_pynew(
    PyTypeObject* type,
    PyObject* args,
    PyObject* kwargs) {
  HANDLE_TH_ERRORS
  TORCH_CHECK(
      type != &THPVariableType,
      "Cannot directly construct TensorBase; subclass it and then construct that");
  jit::tracer::warn("torch.Tensor", jit::tracer::WARN_CONSTRUCTOR);
  // WARNING: tensor is NOT guaranteed to be a fresh tensor; e.g., if it was
  // given a raw pointer that will refcount bump
  // NB: base_tensor_ctor can call into dispatched ATen functions (e.g.,
  // alias(), lift_fresh()) which can return Tensor subclasses.  We allow
  // these to be passed on directly.
  PyObject* obj = THPVariable_WrapWithType(
      torch::utils::base_tensor_ctor(args, kwargs), type);
  if (check_has_torch_dispatch(obj)) {
    THPVariable_Unpack(obj).unsafeGetTensorImpl()->set_python_dispatch(true);
  }
  return obj;
  END_HANDLE_TH_ERRORS
}

static int THPVariable_clear(THPVariable* self) {
  // First clear Tensor specific things
  Py_CLEAR(self->backward_hooks);
  Py_CLEAR(self->post_accumulate_grad_hooks);
  if (self->cdata.defined()) {
    auto pyobj_slot = self->cdata.unsafeGetTensorImpl()->pyobj_slot();
    // Typically the Tensor's pyobj_slot points back to this object. The only
    // time that's not the case is if we had a race in THPVariable_Wrap and we
    // need to discard the Python object because some other thread beat us to
    // setting the pyobj_slot.
    if (pyobj_slot->load_pyobj() == (PyObject*)self) {
      // A Tensor's Python object should only be destroyed when the Tensor has
      // no other references too.
      TORCH_INTERNAL_ASSERT(self->cdata.use_count() == 1);

      // Clear the pyobj_slot so that a try_incref() call from
      // weak_intrusive_ptr::lock() won't see a freed pointer.
      pyobj_slot->clear();
    }
  }
  {
    // MapAllocator can take significant time to release large tensors;
    // release the GIL here to avoid impacting main thread perf.
    pybind11::gil_scoped_release no_gil;
    self->cdata = Variable();
  }
  return 0;
}

static void THPVariable_dealloc(PyObject* self) {
  PyObject_GC_UnTrack(self);
  THPVariable_clear((THPVariable*)self);
  ((THPVariable*)self)->cdata.~Variable();
  Py_TYPE(self)->tp_free(self);
}

static void TORCH_CHECK_TENSOR_SUBTYPE(PyObject* cls) {
  TORCH_CHECK_TYPE(
      PyType_Check(cls),
      "cls must be a type (got ",
      Py_TYPE(cls)->tp_name,
      ")");
  PyTypeObject* type = reinterpret_cast<PyTypeObject*>(cls);
  TORCH_CHECK_TYPE(
      type == &THPVariableType || cls == THPVariableClass ||
          PyType_IsSubtype(type, &THPVariableType),
      "Creating a Tensor subclass from a class that does not inherit from "
      "Tensor is not possible. Make sure your class inherits from Tensor.");
}

/// NOTE [ PyObject Traversal ]
///
/// PyObjects that are wrapping c++ objects can lead to non-trivial traverse
/// logic and it can be tricky to know what to traverse and when. This note
/// tries to clarify what is the danger here and a simple algorithm to choose
/// how to write the tp_traverse and tp_clear functions. If you're not already
/// familiar with how the CPython GC works, you should read this in-depth
/// description: https://devguide.python.org/garbage_collector/
///
/// The complexity for us comes from the fact that some c++ shared_ptr objects
/// own references to python objects and are also owned both by other python
/// objects and c++ objects. This means that to allow the GC to collect all
/// cycles, we need to properly implement the traverse/clear methods that take
/// into account these C++ ownership links.
///
/// The main danger here comes from the fact that, while all python-related code
/// is thread safe wrt the GC execution, other threads might
/// be using our C++ objects arbitrarily which can lead to shared_ptr ref count
/// going up or down in between the different traverse/clear invocations. The
/// one constraint we add here that is not explicitly mentioned in the GC
/// description above is that for a given GC run (meaning while the GIL is
/// held), the traverse/clear pair should never report different ownership
/// relations: if traverse visited a given PyObject, then the clear within that
/// same GC run must still be the sole owner and clear that PyObject.
///
/// A more mechanical algorithm to know what to traverse/clear is as follows:
///   - Any field on this PyObject that contains a strong reference to another
///   PyObject
///     must be visited and cleared. An example of that is the "backward_hooks"
///     field of the THPVariable.
///   - Any field that contains a C++ object that is uniquely owned by this
///   PyObject (either
///     a unique_ptr or a shared_ptr with use_count==1) should have all the
///     PyObject it owns visited and cleared. An example would be here the
///     tensor hooks.
///   - If that uniquely owned C++ object also uniquely owns other C++ objects,
///   these should be
///     visited and cleared as well if they contain any PyObject.
///
/// Caveat: to avoid slow runtime, we limit the depth of this exploration of C++
/// objects in practice and we do not, for example, go through the whole
/// autograd graph, even if it is uniquely owned. This is a known place where
/// users can create noncollectable cycles as described in:
/// https://github.com/pytorch/pytorch/issues/7343
///

static int THPVariable_traverse(PyObject* self, visitproc visit, void* arg) {
  THPVariable* var = reinterpret_cast<THPVariable*>(self);
  Py_VISIT(var->backward_hooks);
  Py_VISIT(var->post_accumulate_grad_hooks);
  const auto& tensor = THPVariable_Unpack(var);
  if (tensor.defined()) {
    // WARNING: The grad_fn traversal logic is very subtle, if you change
    // this, be very careful not to re-introduce this bug:
    // https://gist.github.com/zou3519/7ac92b84dd7d206dcc6eae55fee8372c

    // We ensure that we follow NOTE [ PyObject Traversal ] he by checking
    // that this python object is the sole owner of the underlying Tensor and
    // that this Tensor is the sole owner of its grad_fn. In this case, the
    // only way to get a new reference to the grad_fn is by using this python
    // object, which requires the GIL to be accessed. Note that this is only
    // valid as long as user don't share non-owning references across
    // different threads (which is crazy and should never be done).
    auto autograd_meta = torch::autograd::impl::get_autograd_meta(tensor);
    if (tensor.use_count() == 1) {
      if (autograd_meta) {
        // Do NOT call grad_fn() here as that might trigger a recompute
        const auto& grad_fn = autograd_meta->grad_fn_;
        if (grad_fn && grad_fn.use_count() == 1) {
          // All Node can have a pyobj (stored in "pyobj_")
          Py_VISIT(grad_fn->pyobj());
          // PyNode are special as they also have an "obj" field
          if (auto py_node_fn = dynamic_cast<PyNode*>(grad_fn.get())) {
            Py_VISIT(py_node_fn->obj);
          }
        }
      }
    }
    if (autograd_meta) {
      for (const auto& hook : torch::autograd::impl::hooks(tensor)) {
        if (auto pyhook = dynamic_cast<PyFunctionTensorPreHook*>(hook.get())) {
          Py_VISIT(pyhook->dict);
        }
      }
    }
  }
  return 0;
}

int THPVariableMetaType_init(PyObject* cls, PyObject* args, PyObject* kwargs) {
  if (PyType_Type.tp_init(cls, args, kwargs) < 0) {
    return -1;
  }

  // Don't do anything for the base Tensor class
  if (!THPVariableClass) {
    return 0;
  }

  // Forbid subclassing _TensorBase directly
  py::tuple mro =
      py::reinterpret_borrow<py::tuple>(((PyTypeObject*)cls)->tp_mro);
  bool is_subclass_of_thpvariable = false;
  for (py::handle h : mro) {
    if (h.ptr() == THPVariableClass) {
      is_subclass_of_thpvariable = true;
      break;
    }
  }
  if (!is_subclass_of_thpvariable) {
    PyErr_SetString(PyExc_RuntimeError, "Cannot subclass _TensorBase directly");
    return -1;
  }

  // If the user provided a torch_dispatch implementation, disable
  // torch_function.
  py::object torch_dispatch_impl = py::reinterpret_steal<py::object>(
      PyObject_GetAttrString(cls, "__torch_dispatch__"));
  py::object torch_dispatch_default = py::reinterpret_steal<py::object>(
      PyObject_GetAttrString(THPVariableClass, "__torch_dispatch__"));
  if (torch_dispatch_impl.ptr() != torch_dispatch_default.ptr()) {
    py::object torch_function_impl = py::reinterpret_steal<py::object>(
        PyObject_GetAttrString(cls, "__torch_function__"));
    py::object torch_function_default_bound = py::reinterpret_steal<py::object>(
        PyObject_GetAttrString(THPVariableClass, "__torch_function__"));

    // Since our __torch_function__ is a classmethod, we need to "unbound" the
    // method to get the raw function
    py::object torch_function_default = py::reinterpret_steal<py::object>(
        PyObject_GetAttrString(torch_function_default_bound.ptr(), "__func__"));

    // User-defined __torch_function__ might not be a classmethod
    if (PyObject_HasAttrString(torch_function_impl.ptr(), "__func__")) {
      torch_function_impl = py::reinterpret_steal<py::object>(
          PyObject_GetAttrString(torch_function_impl.ptr(), "__func__"));
    }
    if (torch_function_impl.ptr() == torch_function_default.ptr()) {
      PyObject_SetAttrString(
          cls, "__torch_function__", torch::disabled_torch_function_impl());
    }
  }

  return 0;
}

namespace torch::autograd {

// NOLINTNEXTLINE(modernize-avoid-c-arrays,cppcoreguidelines-avoid-c-arrays,cppcoreguidelines-avoid-non-const-global-variables)
extern PyMethodDef variable_methods[];

static void initTensorImplConversion(PyObject* module) {
  auto m = py::handle(module).cast<py::module>();
  m.def("_wrap_tensor_impl", [](void* ptr) {
    auto p = c10::intrusive_ptr<c10::TensorImpl, at::UndefinedTensorImpl>::
        unsafe_reclaim_from_nonowning(static_cast<c10::TensorImpl*>(ptr));
    TORCH_CHECK(p.defined(), "Can't wrap undefined tensor");
    auto tensor = at::Tensor::wrap_tensor_impl(std::move(p));
    return py::cast(std::move(tensor));
  });
  // set on the module level to avoid mixing pybind and plain CPython extensions
  m.def("_tensor_impl_raw_handle", [](torch::autograd::Variable* t) -> void* {
    // We return a raw non-owning pointer here, we rely on surrounding
    // code to keep the original tensor alive
    return t->getIntrusivePtr().get();
  });
}
} // namespace torch::autograd

bool THPVariable_initModule(PyObject* module) {
  THPVariableMetaType.tp_base = &PyType_Type;
  if (PyType_Ready(&THPVariableMetaType) < 0)
    return false;
  Py_INCREF(&THPVariableMetaType);
  PyModule_AddObject(module, "_TensorMeta", (PyObject*)&THPVariableMetaType);

  static std::vector<PyMethodDef> methods;
  THPUtils_addPyMethodDefs(methods, torch::autograd::variable_methods);
  THPUtils_addPyMethodDefs(methods, extra_methods);
  THPVariableType.tp_methods = methods.data();
  if (PyType_Ready(&THPVariableType) < 0)
    return false;
  Py_INCREF(&THPVariableType);
  PyModule_AddObject(module, "TensorBase", (PyObject*)&THPVariableType);
  Py_INCREF(&THPVariableType);
  PyModule_AddObject(module, "_TensorBase", (PyObject*)&THPVariableType);
#ifdef USE_DISTRIBUTED
  PyModule_AddObject(
      module,
      "__DTensor_fastpath_cache_cleanup",
      py::capsule(
          []() { cleanup_thread_local_native_sharding_propagator_caches(); })
          .release()
          .ptr());
  if (!intern_dtensor_strings()) {
    return false;
  }
  PyModule_AddFunctions(module, extra_dtensor_functions);
#endif
  torch::autograd::initTorchFunctions(module);
  torch::autograd::initTensorImplConversion(module);
  torch::utils::validate_numpy_for_dlpack_deleter_bug();

  return true;
}