Features/746 print0 print toggle

CHANGELOG.md

@@ -15,9 +15,9 @@
 ### Linear Algebra
 - [#842](https://github.com/helmholtz-analytics/heat/pull/842) New feature: `vdot`
 
+## Feature additions
 ### Communication
 - [#868](https://github.com/helmholtz-analytics/heat/pull/868) New `MPICommunication` method `Split`
-
 ### DNDarray
 - [#856](https://github.com/helmholtz-analytics/heat/pull/856) New `DNDarray` method `__torch_proxy__`
 - [#885](https://github.com/helmholtz-analytics/heat/pull/885) New `DNDarray` method `conj`
@@ -31,6 +31,9 @@
 - [#829](https://github.com/helmholtz-analytics/heat/pull/829) New feature: `roll`
 - [#853](https://github.com/helmholtz-analytics/heat/pull/853) New Feature: `swapaxes`
 - [#854](https://github.com/helmholtz-analytics/heat/pull/854) New Feature: `moveaxis`
+### Printing
+- [#816](https://github.com/helmholtz-analytics/heat/pull/816) New feature: Local printing (`ht.local_printing()`) and global printing options
+- [#816](https://github.com/helmholtz-analytics/heat/pull/816) New feature: print only on process 0 with `print0(...)` and `ht.print0(...)`
 ### Random
 - [#858](https://github.com/helmholtz-analytics/heat/pull/858) New Feature: `standard_normal`, `normal`
 ### Rounding

heat/core/printing.py

@@ -1,16 +1,19 @@
 """Allows to output DNDarrays to stdout."""
 
+import builtins
 import copy
 import torch
+from .communication import MPI_WORLD
 
 from .dndarray import DNDarray
 
-__all__ = ["get_printoptions", "set_printoptions"]
+__all__ = ["get_printoptions", "global_printing", "local_printing", "print0", "set_printoptions"]
 
 
 # set the default printing width to a 120
 _DEFAULT_LINEWIDTH = 120
 torch.set_printoptions(profile="default", linewidth=_DEFAULT_LINEWIDTH)
+LOCAL_PRINT = False
 
 # printing
 __PREFIX = "DNDarray"
@@ -24,6 +27,126 @@ def get_printoptions() -> dict:
     return copy.copy(torch._tensor_str.PRINT_OPTS.__dict__)
 
 
+def local_printing() -> None:
+    """
+    The builtin `print` function will now print the local PyTorch Tensor values for
+    `DNDarrays` given as arguments.
+
+    Examples
+    --------
+    >>> x = ht.ht.arange(15 * 5, dtype=ht.float).reshape((15, 5)).resplit(0)
+    >>> ht.local_printing()
+    [0/2]Printing options set to LOCAL. DNDarrays will print the local PyTorch Tensors
+    >>> print(x)
+    [0/2] [[ 0.,  1.,  2.,  3.,  4.],
+    [0/2]  [ 5.,  6.,  7.,  8.,  9.],
+    [0/2]  [10., 11., 12., 13., 14.],
+    [0/2]  [15., 16., 17., 18., 19.],
+    [0/2]  [20., 21., 22., 23., 24.]]
+    [1/2] [[25., 26., 27., 28., 29.],
+    [1/2]  [30., 31., 32., 33., 34.],
+    [1/2]  [35., 36., 37., 38., 39.],
+    [1/2]  [40., 41., 42., 43., 44.],
+    [1/2]  [45., 46., 47., 48., 49.]]
+    [2/2] [[50., 51., 52., 53., 54.],
+    [2/2]  [55., 56., 57., 58., 59.],
+    [2/2]  [60., 61., 62., 63., 64.],
+    [2/2]  [65., 66., 67., 68., 69.],
+    [2/2]  [70., 71., 72., 73., 74.]]
+    """
+    global LOCAL_PRINT
+    LOCAL_PRINT = True
+    print0("Printing options set to LOCAL. DNDarrays will print the local PyTorch Tensors")
+
+
+def global_printing() -> None:
+    """
+    For `DNDarray`s, the builtin `print` function will gather all of the data, format it
+    then print it on ONLY rank 0.
+
+    Returns
+    -------
+    None
+
+    Examples
+    --------
+    >>> x = ht.arange(15 * 5, dtype=ht.float).reshape((15, 5)).resplit(0)
+    >>> print(x)
+    [0] DNDarray([[ 0.,  1.,  2.,  3.,  4.],
+                 [ 5.,  6.,  7.,  8.,  9.],
+                 [10., 11., 12., 13., 14.],
+                 [15., 16., 17., 18., 19.],
+                 [20., 21., 22., 23., 24.],
+                 [25., 26., 27., 28., 29.],
+                 [30., 31., 32., 33., 34.],
+                 [35., 36., 37., 38., 39.],
+                 [40., 41., 42., 43., 44.],
+                 [45., 46., 47., 48., 49.],
+                 [50., 51., 52., 53., 54.],
+                 [55., 56., 57., 58., 59.],
+                 [60., 61., 62., 63., 64.],
+                 [65., 66., 67., 68., 69.],
+                 [70., 71., 72., 73., 74.]], dtype=ht.float32, device=cpu:0, split=0)
+    """
+    global LOCAL_PRINT
+    if not LOCAL_PRINT:
+        return
+    LOCAL_PRINT = False
+    print0(
+        "Printing options set to GLOBAL. DNDarrays will be collected on process 0 before printing"
+    )
+
+
+def print0(*args, **kwargs) -> None:
+    """
+    Wraps the builtin `print` function in such a way that it will only run the command on
+    rank 0. If this is called with DNDarrays and local printing, only the data local to
+    process 0 is printed. For more information see the examples.
+
+    This function is also available as a builtin when importing heat.
+
+    Examples
+    --------
+    >>> x = ht.arange(15 * 5, dtype=ht.float).reshape((15, 5)).resplit(0)
+    >>> # GLOBAL PRINTING
+    >>> ht.print0(x)
+    [0] DNDarray([[ 0.,  1.,  2.,  3.,  4.],
+                 [ 5.,  6.,  7.,  8.,  9.],
+                 [10., 11., 12., 13., 14.],
+                 [15., 16., 17., 18., 19.],
+                 [20., 21., 22., 23., 24.],
+                 [25., 26., 27., 28., 29.],
+                 [30., 31., 32., 33., 34.],
+                 [35., 36., 37., 38., 39.],
+                 [40., 41., 42., 43., 44.],
+                 [45., 46., 47., 48., 49.],
+                 [50., 51., 52., 53., 54.],
+                 [55., 56., 57., 58., 59.],
+                 [60., 61., 62., 63., 64.],
+                 [65., 66., 67., 68., 69.],
+                 [70., 71., 72., 73., 74.]], dtype=ht.float32, device=cpu:0, split=0)
+    >>> ht.local_printing()
+    [0/2] Printing options set to LOCAL. DNDarrays will print the local PyTorch Tensors
+    >>> print0(x)
+    [0/2] [[ 0.,  1.,  2.,  3.,  4.],
+    [0/2]  [ 5.,  6.,  7.,  8.,  9.],
+    [0/2]  [10., 11., 12., 13., 14.],
+    [0/2]  [15., 16., 17., 18., 19.],
+    [0/2]  [20., 21., 22., 23., 24.]], device: cpu:0, split: 0
+    """
+    if not LOCAL_PRINT:
+        args = list(args)
+        for i in range(len(args)):
+            if isinstance(args[i], DNDarray):
+                args[i] = __str__(args[i])
+        args = tuple(args)
+    if MPI_WORLD.rank == 0:
+        print(*args, **kwargs)
+
+
+builtins.print0 = print0
+
+
 def set_printoptions(
     precision=None, threshold=None, edgeitems=None, linewidth=None, profile=None, sci_mode=None
 ):
@@ -67,6 +190,11 @@ def __str__(dndarray) -> str:
     dndarray : DNDarray
         The array for which to obtain the corresponding string
     """
+    if LOCAL_PRINT:
+        return (
+            torch._tensor_str._tensor_str(dndarray.larray, 0)
+            + f", device: {dndarray.device}, split: {dndarray.split}"
+        )
     tensor_string = _tensor_str(dndarray, __INDENT + 1)
     if dndarray.comm.rank != 0:
         return ""

heat/core/tests/test_printing.py

@@ -8,13 +8,49 @@ class TestPrinting(TestCase):
     def setUp(self):
         # move to CPU only for the testing printing, otherwise the compare string will become messy
         ht.use_device("cpu")
+        ht.global_printing()
 
     def tearDown(self):
         # reset the print options back to default after each test run
         ht.set_printoptions(profile="default")
         # reset the default device
         ht.use_device(self.device)
 
+    def test_local_printing(self):
+        x = ht.arange(15 * 5, dtype=ht.float).reshape((15, 5)).resplit(0)
+        global_comp = (
+            "DNDarray([[ 0.,  1.,  2.,  3.,  4.],\n"
+            "          [ 5.,  6.,  7.,  8.,  9.],\n"
+            "          [10., 11., 12., 13., 14.],\n"
+            "          [15., 16., 17., 18., 19.],\n"
+            "          [20., 21., 22., 23., 24.],\n"
+            "          [25., 26., 27., 28., 29.],\n"
+            "          [30., 31., 32., 33., 34.],\n"
+            "          [35., 36., 37., 38., 39.],\n"
+            "          [40., 41., 42., 43., 44.],\n"
+            "          [45., 46., 47., 48., 49.],\n"
+            "          [50., 51., 52., 53., 54.],\n"
+            "          [55., 56., 57., 58., 59.],\n"
+            "          [60., 61., 62., 63., 64.],\n"
+            "          [65., 66., 67., 68., 69.],\n"
+            "          [70., 71., 72., 73., 74.]], dtype=ht.float32, device=cpu:0, split=0)"
+        )
+        if x.comm.rank == 0:
+            self.assertEqual(str(x), global_comp)
+        else:
+            self.assertEqual(str(x), "")
+        ht.local_printing()
+        local_comp = (
+            "[[ 0.,  1.,  2.,  3.,  4.],\n"
+            " [ 5.,  6.,  7.,  8.,  9.],\n"
+            " [10., 11., 12., 13., 14.],\n"
+            " [15., 16., 17., 18., 19.],\n"
+            " [20., 21., 22., 23., 24.]], device: cpu:0, split: 0"
+        )
+        if x.comm.rank == 0 and x.comm.size == 3:
+            self.assertEqual(str(x), local_comp)
+        ht.global_printing()  # needed to keep things correct for the other tests
+
     def test_get_default_options(self):
         print_options = ht.get_printoptions()
         comparison = {