support __format__ for 0-D tensor (#68677)

HydrogenSulfate · web-flow · commit a3a1d3041a34 · 2024-10-14T17:21:21.000+08:00
diff --git a/python/paddle/base/dygraph/tensor_patch_methods.py b/python/paddle/base/dygraph/tensor_patch_methods.py
@@ -913,6 +913,12 @@ def __str__(self: Tensor) -> str:
 
         return tensor_to_string(self)
 
+    def __format__(self, format_spec):
+        if self.ndim == 0:
+            return self.item().__format__(format_spec)
+
+        return object.__format__(self, format_spec)
+
     def __deepcopy__(self, memo: dict[int, Tensor]) -> Tensor:
         """
         Deep copy Tensor, it will always performs Tensor copy.
@@ -1341,6 +1347,7 @@ def __cuda_array_interface__(self):
         ("register_hook", register_hook),
         ("__str__", __str__),
         ("__repr__", __str__),
+        ("__format__", __format__),
         ("__deepcopy__", __deepcopy__),
         ("__module__", "paddle"),
         ("__array__", __array__),
diff --git a/test/legacy_test/test_eager_tensor.py b/test/legacy_test/test_eager_tensor.py
@@ -1280,6 +1280,55 @@ def test___cuda_array_interface__(self):
                     self.assertIn("version", interface)
                     self.assertEqual(interface["version"], 2)
 
+    def test_tensor__format__(self):
+        # test for floating point scalar
+        for width in range(0, 5):
+            paddle_scalar = paddle.randn([])
+            numpy_scalar = paddle_scalar.numpy()
+            format_spec = f".{width}f"
+            self.assertEqual(
+                paddle_scalar.__format__(format_spec),
+                numpy_scalar.__format__(format_spec),
+            )
+            format_spec = f".{width}e"
+            self.assertEqual(
+                paddle_scalar.__format__(format_spec),
+                numpy_scalar.__format__(format_spec),
+            )
+            format_spec = f".{width}g"
+            self.assertEqual(
+                paddle_scalar.__format__(format_spec),
+                numpy_scalar.__format__(format_spec),
+            )
+
+        # test for integer scalar
+        for width in range(0, 5):
+            paddle_scalar = paddle.uniform([], min=-100, max=100).to("int64")
+            numpy_scalar = paddle_scalar.numpy()
+            format_spec = f"{width}d"
+            self.assertEqual(
+                paddle_scalar.__format__(format_spec),
+                numpy_scalar.__format__(format_spec),
+            )
+            format_spec = f"{width}o"
+            self.assertEqual(
+                paddle_scalar.__format__(format_spec),
+                numpy_scalar.__format__(format_spec),
+            )
+            format_spec = f"{width}x"
+            self.assertEqual(
+                paddle_scalar.__format__(format_spec),
+                numpy_scalar.__format__(format_spec),
+            )
+
+        # test for tensor that ndim > 0, expected to raise TypeError
+        paddle_scalar = paddle.uniform([1], min=-100, max=100)
+        self.assertRaises(TypeError, paddle_scalar.__format__, ".3f")
+
+        # test for float scalar but format_spec is 'd', expected to raise ValueError
+        paddle_scalar = paddle.uniform([], min=-100, max=100)
+        self.assertRaises(ValueError, paddle_scalar.__format__, "3d")
+
 
 class TestEagerTensorSetitem(unittest.TestCase):
     def func_setUp(self):
