Update

openxla · Mar 11, 2024 · 75cabb5 · 75cabb5
1 parent 0b90f2b
commit 75cabb5
Showing 1 changed file with 7 additions and 6 deletions.
diff --git a/xla/python/xla_extension/__init__.pyi b/xla/python/xla_extension/__init__.pyi
@@ -696,9 +696,16 @@ class DeviceTopology:
 def buffer_to_dlpack_managed_tensor(
     buffer: ArrayImpl, stream: int | None = None
 ) -> Any: ...
+@overload
 def dlpack_managed_tensor_to_buffer(
     tensor: Any, device: Device, stream: int | None
 ) -> ArrayImpl: ...
+@overload
+def dlpack_managed_tensor_to_buffer( # Legacy overload
+    tensor: Any,
+    cpu_backend: Optional[Client] = ...,
+    gpu_backend: Optional[Client] = ...,
+) -> ArrayImpl: ...
 
 def cuda_array_interface_to_buffer(
     cai: Dict[str, Union[
@@ -710,12 +717,6 @@ def cuda_array_interface_to_buffer(
     gpu_backend: Optional[Client] = ...,
 ) -> ArrayImpl: ...
 
-# Legacy overload
-def dlpack_managed_tensor_to_buffer(
-    tensor: Any,
-    cpu_backend: Optional[Client] = ...,
-    gpu_backend: Optional[Client] = ...,
-) -> ArrayImpl: ...
 
 # === BEGIN py_traceback.cc