Method __dlpack__

torch/_tensor.py:1426–1482

Creates a DLpack `capsule https://data-apis.org/array-api/latest/design_topics/data_interchange.html#data-interchange`_ of the current tensor to be exported to other libraries. This function will be called from the `from_dlpack` method of the library that will consu

(self, stream=None)

    __torch_dispatch__ = _C._disabled_torch_dispatch_impl

    def __dlpack__(self, stream=None):
        """
        Creates a DLpack `capsule https://data-apis.org/array-api/latest/design_topics/data_interchange.html#data-interchange`_
        of the current tensor to be exported to other libraries.

        This function will be called from the `from_dlpack` method
        of the library that will consume the capsule. `from_dlpack` passes the current
        stream to this method as part of the specification.

        Args:
            stream (integer or None): An optional Python integer representing a
            pointer to a CUDA stream. The current stream is synchronized with
            this stream before the capsule is created, and since the capsule
            shares its storage with the tensor this make it safe to access from
            both streams.  If None or -1 is passed then no synchronization is performed.
            If 1 (on CUDA) or 0 (on ROCM) then the default stream is used for
            synchronization.
        """
        if has_torch_function_unary(self):
            return handle_torch_function(Tensor.__dlpack__, (self,), self, stream)

        # DLPack capsules can't capture all of PyTorch's semantics,
        # so we prohibit exporting tensors that would lose their properties like
        # requires_grad and having the conjugate bit set.
        if self.requires_grad:
            raise RuntimeError(
                "Can't export tensors that require gradient, use tensor.detach()"
            )
        if self.is_conj():
            raise RuntimeError("Can't export tensors with the conjugate bit set")
        if self.layout != torch.strided:
            raise RuntimeError(
                "Can't export tensors with layout other than torch.strided"
            )

        if stream is not None and type(stream) is not int:
            # Stream pointers in CUDA/ROCm are uniquely numbered and can
            # be retrieved from their integer value.
            raise TypeError("stream must be ``int`` or ``none``")
        elif stream is not None and stream != -1:
            if self.device.type == "cuda":
                # NB: This logic handles the special case values for default
                # streams and must be kept in sync with from_dlpack in
                # torch/utils/dlpack.py
                if stream == 1 and torch.version.hip is None:
                    stream = torch.cuda.default_stream()
                elif stream == 0 and torch.version.hip is not None:
                    stream = torch.cuda.default_stream()
                else:
                    stream = torch.cuda.ExternalStream(stream)
                # Only synchronize on different streams
                sync_stream = torch.cuda.current_stream()
                if stream != sync_stream:
                    event = torch.cuda.Event()
                    event.record(sync_stream)
                    stream.wait_event(event)
        return torch.to_dlpack(self)