Return torch tensors from the Python runtime only when torch is imported

docs/source/runtime-python-api-reference.rst

@@ -26,7 +26,7 @@ For detailed information on how APIs evolve and the deprecation process, please
 Without PyTorch
 ~~~~~~~~~~~~~~~
 
-The runtime also works when PyTorch is not installed. Then ``Method.execute`` takes NumPy arrays as inputs and returns ``ExecuTorchResult`` values, which NumPy can read. For a model that adds its two inputs:
+ExecuTorch never imports PyTorch on its own. ``Method.execute`` returns ``torch.Tensor`` outputs only when your program has already imported ``torch``, and ``ExecuTorchResult`` values otherwise, even when PyTorch is installed. So a program that never imports ``torch`` runs without loading it. It passes NumPy arrays as inputs and reads the ``ExecuTorchResult`` outputs with NumPy. For a model that adds its two inputs:
 
 .. code-block:: python
 
@@ -43,7 +43,7 @@ Other buffers, such as ``memoryview``, work as inputs too. NumPy has no bfloat16
 
 For data on a GPU, use DLPack, the standard way to share tensors between Python libraries. This needs ExecuTorch built with CUDA support. Pass a CuPy array to a method that takes CUDA inputs. Read a CUDA result with ``cupy.from_dlpack(out)``, or copy it to the CPU with ``out.to_host()``. ``ExecuTorchResult.from_host()`` copies CPU data to a CUDA device.
 
-To get ``ExecuTorchResult`` values even when PyTorch is installed, set ``EXECUTORCH_PYBINDINGS_TENSOR_OUTPUT=executorch`` before importing ExecuTorch.
+To always get ``ExecuTorchResult`` values, even after ``torch`` is imported, set ``EXECUTORCH_PYBINDINGS_TENSOR_OUTPUT=executorch`` before importing ExecuTorch. Set it to ``torch`` to always get ``torch.Tensor`` outputs; ExecuTorch then imports PyTorch itself.
 
 CUDA streams
 ~~~~~~~~~~~~

examples/arm/person_detection_example_cortex_m/utils/pte_metadata.py

@@ -24,6 +24,8 @@ def io_method_names() -> tuple[str, ...]:
 
 
 def read_io_qparams(path: Path) -> dict[str, float | int]:
+    # Tensor outputs have item() only as torch tensors, once torch is imported.
+    import torch  # noqa: F401
     from executorch.runtime import Runtime
 
     program = Runtime.get().load_program(path)

extension/pybindings/README.md

@@ -91,9 +91,12 @@ with `clone_outputs=False` lives in the method's output memory, so passing it
 back in as an input means each run reads what the run before it wrote there;
 pass a copied result instead.
 
-To select `ExecuTorchResult` outputs even when PyTorch is installed, set the
-environment variable before importing the bindings. The setting is fixed when
-the extension loads, so importing PyTorch later does not change the output type.
+The bindings never import PyTorch on their own. Each call returns torch tensors
+only if `torch` is already imported, and `ExecuTorchResult` values otherwise, so
+a program that never imports `torch` does not load it. To fix the output type,
+set the environment variable before importing the bindings: `executorch` keeps
+`ExecuTorchResult` outputs even after `torch` is imported, and `torch` makes the
+bindings import PyTorch and always return torch tensors.
 
 ```python
 import os
@@ -103,11 +106,9 @@ from executorch.extension.pybindings import portable_lib
 ```
 
 `ExecuTorchResult` preserves shape, strides, dtype, and device. CPU results expose
-the Python buffer protocol so NumPy can read them without another copy. Without
-the environment setting, the bindings fall back to `ExecuTorchResult` if
-PyTorch cannot be imported. Explicitly selecting `torch` requires it to import
-successfully. ATen-mode bindings continue to return native PyTorch tensors
-directly.
+the Python buffer protocol so NumPy can read them without another copy.
+Explicitly selecting `torch` requires it to import successfully. ATen-mode
+bindings continue to return native PyTorch tensors directly.
 
 ```python
 output = module.forward([numpy_input])[0]

extension/pybindings/portable_lib.py

@@ -28,23 +28,19 @@ class _LightweightExperimentalWarning(DeprecationWarning):
 
 _TENSOR_OUTPUT_ENV = "EXECUTORCH_PYBINDINGS_TENSOR_OUTPUT"
 _tensor_output = os.environ.get(_TENSOR_OUTPUT_ENV)
-if _tensor_output in (None, "torch"):
+# Without the setting, ExecuTorch never imports torch itself: each call returns
+# torch tensors only if the caller has already imported torch.
+if _tensor_output == "torch":
     try:
-        # Preserve the existing API: importing portable_lib also imports Torch,
-        # and tensor outputs are torch.Tensor objects.
-        import torch as _torch
+        _import_module("torch")
     except ImportError as error:
-        if _tensor_output == "torch":
-            raise ImportError(
-                "ExecuTorch portable pybindings were configured to return "
-                "torch.Tensor outputs, but PyTorch could not be imported. "
-                f"Install PyTorch or set {_TENSOR_OUTPUT_ENV}=executorch "
-                "before importing portable_lib."
-            ) from error
-        _torch = None
-elif _tensor_output == "executorch":
-    _torch = None
-else:
+        raise ImportError(
+            "ExecuTorch portable pybindings were configured to return "
+            "torch.Tensor outputs, but PyTorch could not be imported. "
+            f"Install PyTorch or set {_TENSOR_OUTPUT_ENV}=executorch "
+            "before importing portable_lib."
+        ) from error
+elif _tensor_output not in (None, "executorch"):
     raise ValueError(
         f"{_TENSOR_OUTPUT_ENV} must be either 'torch' or 'executorch', "
         f"but was {_tensor_output!r}"
@@ -152,7 +148,6 @@ class _LightweightExperimentalWarning(DeprecationWarning):
 del _bindings
 del _tensor_output
 del _TENSOR_OUTPUT_ENV
-del _torch
 del _warning_category
 del _import_module
 del _exir_warnings

extension/pybindings/pybindings.cpp

@@ -627,38 +627,43 @@ struct RuntimeErrorScope final {
   std::unique_ptr<PendingLogs> logs_;
 };
 
-#ifndef USE_ATEN_LIB
-constexpr char kTensorOutputEnvironmentVariable[] =
-    "EXECUTORCH_PYBINDINGS_TENSOR_OUTPUT";
-
 enum class PortableTensorOutput { Torch, ExecuTorch };
 
-PortableTensorOutput& portable_tensor_output() {
-  static PortableTensorOutput output = PortableTensorOutput::Torch;
+std::optional<PortableTensorOutput>& tensor_output_override() {
+  static std::optional<PortableTensorOutput> output;
   return output;
 }
 
+// Decided for each call, so a caller that never imports torch never gets it
+// imported by these bindings and receives ExecuTorchResult outputs.
+PortableTensorOutput portable_tensor_output() {
+  if (const auto& output = tensor_output_override()) {
+    return *output;
+  }
+  return loaded_torch_module() != nullptr ? PortableTensorOutput::Torch
+                                          : PortableTensorOutput::ExecuTorch;
+}
+
+#ifndef USE_ATEN_LIB
+constexpr char kTensorOutputEnvironmentVariable[] =
+    "EXECUTORCH_PYBINDINGS_TENSOR_OUTPUT";
+
 void configure_portable_tensor_output() {
   const char* value = std::getenv(kTensorOutputEnvironmentVariable);
-  if (value != nullptr && std::strcmp(value, "executorch") == 0) {
-    portable_tensor_output() = PortableTensorOutput::ExecuTorch;
+  if (value == nullptr) {
+    return;
+  }
+  if (std::strcmp(value, "executorch") == 0) {
+    tensor_output_override() = PortableTensorOutput::ExecuTorch;
     return;
   }
-  if (value != nullptr && std::strcmp(value, "torch") != 0) {
+  if (std::strcmp(value, "torch") != 0) {
     throw py::value_error(
         std::string(kTensorOutputEnvironmentVariable) +
         " must be either 'torch' or 'executorch', but was '" + value + "'");
   }
-  try {
-    py::module_::import("torch");
-    (void)loaded_torch_module();
-    portable_tensor_output() = PortableTensorOutput::Torch;
-  } catch (py::error_already_set& error) {
-    if (value != nullptr || !error.matches(PyExc_ImportError)) {
-      throw;
-    }
-    portable_tensor_output() = PortableTensorOutput::ExecuTorch;
-  }
+  py::module_::import("torch");
+  tensor_output_override() = PortableTensorOutput::Torch;
 }
 #endif
 
@@ -1978,9 +1983,9 @@ py::object dlpack_output_owner(
 py::object module_aliased_torch_output(
     const executorch::aten::Tensor& tensor,
     const std::vector<std::shared_ptr<TorchTensorView>>& torch_inputs,
-    const std::string& method_name) {
-  if (torch_inputs.empty() ||
-      portable_tensor_output() != PortableTensorOutput::Torch ||
+    const std::string& method_name,
+    PortableTensorOutput output) {
+  if (torch_inputs.empty() || output != PortableTensorOutput::Torch ||
       tensor.nbytes() == 0) {
     return py::none();
   }
@@ -2414,9 +2419,10 @@ class PyCallerStreamScope final {
 
 py::object portable_tensor_result(
     const executorch::aten::Tensor& tensor,
+    PortableTensorOutput output,
     bool clone = true,
     py::object owner = py::none()) {
-  if (portable_tensor_output() == PortableTensorOutput::ExecuTorch) {
+  if (output == PortableTensorOutput::ExecuTorch) {
     return py::cast(std::make_shared<PyExecuTorchResult>(
         tensor, clone, clone ? py::none() : std::move(owner)));
   }
@@ -2661,7 +2667,8 @@ inline py::list get_outputs_as_py_list(
     bool clone_outputs = true,
     py::object owner = py::none(),
     const std::vector<py::object>& output_owners = {},
-    const std::vector<py::object>& handed_over = {}) {
+    const std::vector<py::object>& handed_over = {},
+    PortableTensorOutput output = portable_tensor_output()) {
   const auto outputs_size = outputs.size();
   py::list list(outputs_size);
   for (size_t i = 0; i < outputs_size; ++i) {
@@ -2690,9 +2697,10 @@ inline py::list get_outputs_as_py_list(
           i < handed_over.size() && handed_over[i];
       list[i] = returns_without_copy
           ? portable_tensor_result(
-                v.toTensor(), /*clone=*/false, handed_over[i])
+                v.toTensor(), output, /*clone=*/false, handed_over[i])
           : portable_tensor_result(
                 v.toTensor(),
+                output,
                 clone_outputs,
                 output_owners.empty() ? owner : output_owners[i]);
 #endif
@@ -3538,20 +3546,22 @@ struct PyModule final {
                             : py::none());
     }
 #endif
+    const auto tensor_output = portable_tensor_output();
     auto result = get_outputs_as_py_list(
         outputs.get(),
         clone_outputs,
         std::move(output_owner),
         output_owners,
         clone_outputs ? handed_over_device_outputs(outputs.get())
-                      : std::vector<py::object>());
+                      : std::vector<py::object>(),
+        tensor_output);
 #ifndef USE_ATEN_LIB
     if (!clone_outputs) {
       for (size_t i = 0; i < outputs->size(); ++i) {
         const auto& output = (*outputs)[i];
         if (output.isTensor()) {
           auto alias = module_aliased_torch_output(
-              output.toTensor(), torch_inputs, method_name);
+              output.toTensor(), torch_inputs, method_name, tensor_output);
           if (!alias.is_none()) {
             result[i] = std::move(alias);
           }
@@ -4568,7 +4578,8 @@ struct PyMethod final {
     }
 
     // Retrieve outputs
-    py::list outputs = get_outputs_as_py_list(result, clone_outputs);
+    py::list outputs =
+        get_outputs_as_py_list(result, clone_outputs, portable_tensor_output());
     for (size_t i = 0; i < caller_backed.size(); ++i) {
       if (caller_backed[i]) {
         outputs[i] = caller_backed[i];
@@ -4606,11 +4617,13 @@ struct PyMethod final {
     // writable: a torch tensor cannot be read-only, so torch gets a copy.
     const bool constant = attr->nbytes() != 0 &&
         !storage_->memory->owns_data(attr->const_data_ptr());
-    if (constant && portable_tensor_output() == PortableTensorOutput::Torch) {
-      return portable_tensor_result(attr.get(), /*clone=*/true);
+    const auto output = portable_tensor_output();
+    if (constant && output == PortableTensorOutput::Torch) {
+      return portable_tensor_result(attr.get(), output, /*clone=*/true);
     }
     auto result = portable_tensor_result(
         attr.get(),
+        output,
         /*clone=*/false,
         borrowed_output_owner(attr.get()));
     if (constant) {
@@ -4754,9 +4767,10 @@ struct PyMethod final {
     std::vector<std::shared_ptr<DLPackTensorView>> dlpack_inputs;
   };
 
-  py::object aliased_torch_output(const executorch::aten::Tensor& tensor) {
-    if (torch_inputs_.empty() ||
-        portable_tensor_output() != PortableTensorOutput::Torch ||
+  py::object aliased_torch_output(
+      const executorch::aten::Tensor& tensor,
+      PortableTensorOutput output) {
+    if (torch_inputs_.empty() || output != PortableTensorOutput::Torch ||
         tensor.nbytes() == 0) {
       return py::none();
     }
@@ -4808,7 +4822,8 @@ struct PyMethod final {
   // result stays independent.
   std::optional<py::object> hand_over_device_output(
       size_t index,
-      const executorch::aten::Tensor& tensor) {
+      const executorch::aten::Tensor& tensor,
+      PortableTensorOutput output) {
     const auto slot = device_buffer_slot(tensor);
     if (!slot) {
       return std::nullopt;
@@ -4824,10 +4839,11 @@ struct PyMethod final {
         std::make_shared<std::shared_ptr<DeviceMemoryBuffer>>(buffer);
     const std::shared_ptr<void> memory(holder, buffer->data());
     py::object result;
-    if (portable_tensor_output() == PortableTensorOutput::Torch) {
+    if (output == PortableTensorOutput::Torch) {
       auto* retained = new std::shared_ptr<void>(memory);
       result = portable_tensor_result(
           tensor,
+          output,
           /*clone=*/false,
           py::capsule(retained, [](void* pointer) {
             delete static_cast<std::shared_ptr<void>*>(pointer);
@@ -5008,7 +5024,8 @@ struct PyMethod final {
 
   py::list get_outputs_as_py_list(
       const std::vector<EValue>& outputs,
-      bool clone_outputs = true) {
+      bool clone_outputs,
+      PortableTensorOutput output) {
     const auto outputs_size = outputs.size();
     py::list list(outputs_size);
     for (size_t i = 0; i < outputs_size; ++i) {
@@ -5042,20 +5059,22 @@ struct PyMethod final {
           continue;
         }
         if (clone_outputs) {
-          if (const auto handed = hand_over_device_output(i, v.toTensor())) {
+          if (const auto handed =
+                  hand_over_device_output(i, v.toTensor(), output)) {
             list[i] = *handed;
             continue;
           }
         }
         if (!clone_outputs) {
-          auto alias = aliased_torch_output(v.toTensor());
+          auto alias = aliased_torch_output(v.toTensor(), output);
           if (!alias.is_none()) {
             list[i] = std::move(alias);
             continue;
           }
         }
         list[i] = portable_tensor_result(
             v.toTensor(),
+            output,
             clone_outputs,
             clone_outputs && v.toTensor().device().is_cpu()
                 ? py::none()
@@ -5407,9 +5426,13 @@ PYBIND11_MODULE(EXECUTORCH_PYTHON_MODULE_NAME, m) {
 #else
   configure_portable_tensor_output();
   m.attr("_uses_aten") = false;
-  m.attr("_tensor_output") =
-      portable_tensor_output() == PortableTensorOutput::Torch ? "torch"
-                                                              : "executorch";
+  // None means each call picks, depending on whether torch is imported.
+  const auto& tensor_output = tensor_output_override();
+  m.attr("_tensor_output") = !tensor_output
+      ? py::object(py::none())
+      : py::str(
+            *tensor_output == PortableTensorOutput::Torch ? "torch"
+                                                          : "executorch");
 #endif
   py::class_<PyExecuTorchResult, std::shared_ptr<PyExecuTorchResult>>(
       m,

extension/pybindings/pybindings_executorch_result.h

@@ -97,10 +97,17 @@ inline PyObject* loaded_torch_module() {
   static PyObject* torch_module = nullptr;
   if (torch_module == nullptr) {
     PyObject* loaded = PyDict_GetItemString(PyImport_GetModuleDict(), "torch");
-    if (loaded != nullptr) {
-      Py_INCREF(loaded);
-      torch_module = loaded;
+    if (loaded == nullptr || loaded == Py_None) {
+      return nullptr;
     }
+    // A torch that another thread is still importing is not usable yet.
+    const auto spec = py::getattr(py::handle(loaded), "__spec__", py::none());
+    if (!spec.is_none() &&
+        py::getattr(spec, "_initializing", py::bool_(false)).cast<bool>()) {
+      return nullptr;
+    }
+    Py_INCREF(loaded);
+    torch_module = loaded;
   }
   return torch_module;
 }