Method empty

python/tvm/runtime/disco/session.py:130–167

Create an empty Tensor on all workers and attach them to a DRef. Parameters ---------- shape : tuple of int The shape of the Tensor. dtype : str The data type of the Tensor. device : Optional[Device] = None The device of

(
        self,
        shape: Sequence[int],
        dtype: str,
        device: Device | None = None,
        worker0_only: bool = False,
        in_group: bool = True,
    )

        return func

    def empty(
        self,
        shape: Sequence[int],
        dtype: str,
        device: Device | None = None,
        worker0_only: bool = False,
        in_group: bool = True,
    ) -> DRef:
        """Create an empty Tensor on all workers and attach them to a DRef.

        Parameters
        ----------
        shape : tuple of int
            The shape of the Tensor.

        dtype : str
            The data type of the Tensor.

        device : Optional[Device] = None
            The device of the Tensor.

        worker0_only: bool
            If False (default), allocate an array on each worker.  If
            True, only allocate an array on worker0.

        in_group: bool
            Take effective when `worker0_only` is True. If True (default),
            allocate an array on each first worker in each group. If
            False, only allocate an array on worker0 globally.

        Returns
        -------
        array : DRef
            The created Tensor.

        """
        func = self._get_cached_method("runtime.disco.empty")
        return func(Shape(shape), dtype, device, worker0_only, in_group)

    def shutdown(self):
        """Shut down the Disco session"""