dtensor_from_fn¶

paddle.distributed. dtensor_from_fn ( fn, mesh, placements, *args, **kwargs ) [source]

Construct a Distributed Tensor from a function of arguments.

Parameters

fn (callable) – A callable function that takes arguments of Distributed Tensor and returns 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.
*args (tuple) – A tuple of arguments to be passed to the fn function.
**kwargs (dict) – A dict of arguments to be passed to the fn function.

Returns

A Tensor constructed from fn with distributed attributes.

Return type

Tensor

Examples

           >>> import paddle
>>> import paddle.distributed as dist
>>> # Create a distributed attribute
>>> mesh = dist.ProcessMesh([0, 1], dim_names=["x"])
>>> # Call the function dtensor_from_fn with dist_attr parameter
>>> d_tensor = dist.dtensor_from_fn(paddle.ones, mesh, [dist.Replicate()], shape=[1])
>>> print(d_tensor)

          