[AutoParallel] Add paddle.distributed.dtensor_from_fn api

python/paddle/distributed/__init__.py

@@ -66,6 +66,7 @@
 
 from .auto_parallel import shard_op  # noqa: F401
 from .auto_parallel.api import shard_tensor  # noqa: F401
+from .auto_parallel.api import dtensor_from_fn  # noqa: F401
 
 from .fleet import BoxPSDataset  # noqa: F401
 
@@ -126,4 +127,5 @@
     "ProcessMesh",
     "DistAttr",
     "shard_tensor",
+    "dtensor_from_fn",
 ]

python/paddle/distributed/auto_parallel/api.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+
 import paddle
 from paddle.distributed.auto_parallel.interface import (
     shard_tensor as shard_tensor_static,
@@ -69,19 +70,6 @@ def __init__(self, mesh, sharding_specs):
         self.process_mesh = mesh
         self.dims_mapping = dims_mapping
 
-        self.mark_annotated("process_mesh")
-        self.mark_annotated("dims_mapping")
-
-    @property
-    def sharding_specs(self):
-        """
-        Get sharding_specs of the dist_attr
-
-        Returns:
-            list[str]: sharding_specs
-        """
-        return self._sharding_specs
-
 
 def shard_tensor(
     data, dtype=None, place=None, stop_gradient=True, dist_attr=None
@@ -142,3 +130,36 @@ def shard_tensor(
         return shard_tensor_static(
             data, dist_attr.process_mesh, dist_attr.sharding_specs
         )
+
+
+def dtensor_from_fn(fn, dist_attr, *args, **kwargs):
+    """
+    Construct a Distributed Tensor from a paddle api function of arguments.
+
+    Args:
+        fn (callable): A paddle api function that takes arguments of *args, **kwargs and returns tensor.
+        dist_attr(paddle.distributed.DistAttr): Specify how tensors are distributed or sliced on ProcessMesh.
+        *args: A list of arguments to be passed to the ``fn`` function.
+        **kwargs: A list of arguments to be passed to the ``fn`` function.
+
+    Retruns:
+        Tensor: A Tensor constructed from ``fn`` with distributed attributes.
+
+    Examples:
+
+    .. code-block:: python
+
+    >>> import paddle
+    >>> import paddle.distribute as dist
+
+    >>> # Create a distributed attribute
+    >>> mesh = dist.ProcessMesh([[2, 4, 5], [0, 1, 3]], dim_names=["x", "y"])
+    >>> dist_attr = dist.DistAttr(mesh=mesh, sharding_specs=['x', 'y'])
+
+    >>> # Call the function dtensor_from_fn with dist_attr parameter
+    >>> d_tensor = dist.dtensor_from_fn(paddle.ones, dist_attr=dist_attr, shape=[2, 3])
+
+    >>> print(d_tensor)
+    """
+    tensor = fn(*args, **kwargs)
+    return shard_tensor(tensor, dist_attr=dist_attr)

test/auto_parallel/test_dist_tensor.py

@@ -53,6 +53,56 @@ def test_dist_tensor_creation(self):
         self.assertEqual(dist_tensor_with_tensor.dist_attr, dist_attr)
 
 
+class TestDistributedTensor(unittest.TestCase):
+    def test_dtensor_from_fn(self):
+        # Create a distributed attribute
+        mesh = dist.ProcessMesh([[2, 4, 5], [0, 1, 3]], dim_names=["x", "y"])
+        dist_attr = dist.DistAttr(mesh=mesh, sharding_specs=['x', 'y'])
+
+        # Test with generate_tensor_ones()
+        # Call the function dtensor_from_fn with dist_attr parameter
+        result = dist.dtensor_from_fn(
+            paddle.ones, dist_attr=dist_attr, shape=[2, 3]
+        )
+
+        # Verify the result
+        self.assertIsInstance(result, paddle.Tensor)
+        self.assertEqual(result.shape, [2, 3])
+        self.assertEqual(result.dist_attr, dist_attr)
+
+        # Test with generate_tensor_zeros()
+        result_zeros = dist.dtensor_from_fn(
+            paddle.zeros, dist_attr=dist_attr, shape=[2, 3]
+        )
+        self.assertIsInstance(result_zeros, paddle.Tensor)
+        self.assertEqual(result_zeros.shape, [2, 3])
+        self.assertEqual(result_zeros.dist_attr, dist_attr)
+
+        # Test with generate_tensor_random()
+        result_random = dist.dtensor_from_fn(
+            paddle.rand, dist_attr=dist_attr, shape=[2, 3]
+        )
+        self.assertIsInstance(result_random, paddle.Tensor)
+        self.assertEqual(result_random.shape, [2, 3])
+        self.assertEqual(result_random.dist_attr, dist_attr)
+
+        # Test with invalid sharding_specs length
+        with self.assertRaises(AssertionError):
+            invalid_dist_attr = dist.DistAttr(mesh=mesh, sharding_specs=['x'])
+            dist.dtensor_from_fn(
+                paddle.ones, dist_attr=invalid_dist_attr, shape=[2, 3]
+            )
+
+        # Test exceptions when running in static mode
+        paddle.enable_static()
+        with self.assertRaises(NotImplementedError):
+            with paddle.static.program_guard(paddle.static.Program()):
+                dist.dtensor_from_fn(
+                    paddle.ones, dist_attr=dist_attr, shape=[2, 3]
+                )
+        paddle.disable_static()
+
+
 class TestDistTensorForDygraphAPI(unittest.TestCase):
     def check_tensor_eq(self, a, b):
         np1 = a.numpy()