Function shard_tensor

python/paddle/distributed/auto_parallel/api.py:246–393

Creates a distributed Tensor (i.e., Tensor with distributed attributes or DistTensor for short) from the input data, which can be a scalar, tuple, list, numpy.ndarray, or paddle.Tensor. If the ``data`` is already a Tensor, it will be transformed into a distributed Tensor. Args:

(
    data: Tensor | TensorLike | NestedNumericSequence,
    mesh: ProcessMesh,
    placements: Sequence[Placement],
    dtype: DTypeLike | None = None,
    place: PlaceLike | None = None,
    stop_gradient: bool | None = None,
)

def shard_tensor(
    data: Tensor | TensorLike | NestedNumericSequence,
    mesh: ProcessMesh,
    placements: Sequence[Placement],
    dtype: DTypeLike | None = None,
    place: PlaceLike | None = None,
    stop_gradient: bool | None = None,
) -> Tensor:
    """
    Creates a distributed Tensor (i.e., Tensor with distributed attributes or DistTensor for short)
    from the input data, which can be a scalar, tuple, list, numpy.ndarray, or paddle.Tensor.

    If the ``data`` is already a Tensor, it will be transformed into a distributed Tensor.

    Args:
        data(scalar|tuple|list|ndarray|Tensor): Initial data for the tensor.
            Can be a scalar, list, tuple, numpy.ndarray, paddle.Tensor.
        mesh(paddle.distributed.ProcessMesh): The `ProcessMesh` object describes the Cartesian topology of the used processes.
        placements(list[paddle.distributed.Placement]): the placements describe how to place the tensor on ProcessMesh, it can
            be Shard, Replicate and Partial.
        dtype(str|paddle.dtype|np.dtype, optional): The desired data type of returned tensor.
            It Can be 'bool' , 'float16' , 'float32' , 'float64' , 'int8' , 'int16' , 'int32' , 'int64' , 'uint8',
            'complex64' , 'complex128'. Default: None. If None, the the dtype is inferred from ``data``
            except for python float number, in which case the dtype is inferred from ``get_default_type`` .
        place(CPUPlace|CUDAPinnedPlace|CUDAPlace|str, optional): The place to allocate Tensor. Can be
            CPUPlace, CUDAPinnedPlace, CUDAPlace. Default: None, means global place. If ``place`` is
            string, It can be ``cpu``, ``gpu:x`` and ``gpu_pinned``, where ``x`` is the index of the GPUs.
        stop_gradient(bool, optional): Whether to block the gradient propagation of Autograd. If
            ``stop_gradient`` is None, set the returned Tensor's ``stop_gradient`` identical as the
            ``data.stop_gradient`` when ``data`` has ``stop_gradient`` attribute and True otherwise.
            Default: None.

    Returns:
        Tensor: A Tensor constructed from ``data`` with distributed attributes.

    Examples:
        .. code-block:: pycon

            >>> import paddle
            >>> import paddle.distributed as dist

            >>> mesh = dist.ProcessMesh([[2, 4, 5], [0, 1, 3]], dim_names=['x', 'y'])

            >>> # dense tensor
            >>> a = paddle.to_tensor(
            ...     [
            ...         [1, 2, 3],
            ...         [5, 6, 7],
            ...     ]
            ... )

            >>> # doctest: +REQUIRES(env:DISTRIBUTED)
            >>> # distributed tensor
            >>> d_tensor = dist.shard_tensor(a, mesh, [dist.Shard(0), dist.Shard(1)])

            >>> print(d_tensor)

    """