[Feature] Add set_printoptions for configurable TensorDict repr

Add set_printoptions / get_printoptions to control which attributes
appear in TensorDict's __repr__.  Works as a global setter, context
manager, or decorator.

Options: show/hide batch_size, device, is_shared at the TensorDict
level; show/hide shape, device, dtype, is_shared at the per-tensor
level; opt-in extended attributes (requires_grad, is_contiguous,
is_view, storage_size, plain value summary).

Also adds a "Printing and Display" documentation page covering the
feature, the motivation behind TensorDict's metadata-first repr, and
parse_tensor_dict_string.

Made-with: Cursor

Author: vmoens

docs/source/index.rst

@@ -108,6 +108,7 @@ Contents
    distributed
    fx
    saving
+   printing
    reference/index
 
 Indices and tables

docs/source/printing.rst

@@ -0,0 +1,166 @@
+Printing and Display
+====================
+
+Why printing a TensorDict is more useful than printing a tensor
+---------------------------------------------------------------
+
+When working with PyTorch tensors, calling ``print(tensor)`` dumps the raw
+numerical content.  In practice, however, most ``print`` calls during debugging
+are motivated by a single question: *what does this tensor look like?*  You want
+its shape, dtype, device, maybe whether it requires a gradient -- not a wall of
+floating-point numbers.
+
+Because a :class:`~tensordict.TensorDict` groups multiple tensors under named
+keys, its ``__repr__`` gives you exactly that -- a structured, at-a-glance
+summary of every tensor it contains:
+
+    >>> import torch
+    >>> from tensordict import TensorDict
+    >>> td = TensorDict(
+    ...     image=torch.randn(32, 3, 64, 64),
+    ...     label=torch.randint(10, (32,)),
+    ...     batch_size=[32],
+    ... )
+    >>> print(td)
+    TensorDict(
+        fields={
+            image: Tensor(shape=torch.Size([32, 3, 64, 64]), device=cpu, dtype=torch.float32, is_shared=False),
+            label: Tensor(shape=torch.Size([32]), device=cpu, dtype=torch.int64, is_shared=False)},
+        batch_size=torch.Size([32]),
+        device=None,
+        is_shared=False)
+
+No data is printed, no truncation ellipses, no guessing at dimensionality.
+One glance tells you the names, shapes, dtypes and devices of everything in
+the batch.
+
+Configuring the display with ``set_printoptions``
+-------------------------------------------------
+
+By default, every attribute is shown for backward compatibility.  In many
+situations, though, some of those attributes are noise.  For instance, if all
+your work is on CPU and nothing is shared, ``device=cpu`` and
+``is_shared=False`` are repeated on every line without adding information.
+
+:class:`~tensordict.set_printoptions` lets you control exactly which attributes
+appear.  It works as a **global setter**, a **context manager** or a
+**decorator**, following the same pattern as :class:`~tensordict.set_lazy_legacy`
+and :func:`torch.set_printoptions`.
+
+Global configuration
+~~~~~~~~~~~~~~~~~~~~
+
+Call :meth:`~tensordict.set_printoptions.set` to change the defaults for the
+rest of the process:
+
+    >>> from tensordict import set_printoptions
+    >>> set_printoptions(show_device=False, show_is_shared=False).set()
+    >>> print(td)
+    TensorDict(
+        fields={
+            image: Tensor(shape=torch.Size([32, 3, 64, 64]), dtype=torch.float32),
+            label: Tensor(shape=torch.Size([32]), dtype=torch.int64)},
+        batch_size=torch.Size([32]))
+
+Scoped configuration (context manager)
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+Use the context-manager form when you only want the change for a specific
+block of code.  The previous settings are automatically restored on exit:
+
+    >>> from tensordict import set_printoptions
+    >>> with set_printoptions(show_tensor_dtype=False, show_is_shared=False):
+    ...     print(td)  # dtype and is_shared hidden
+    >>> print(td)  # back to defaults
+
+Decorator
+~~~~~~~~~
+
+You can also decorate a function so that every ``repr`` call inside it uses
+the specified options:
+
+    >>> @set_printoptions(show_is_shared=False)
+    ... def summarise(td):
+    ...     print(td)
+
+Available options
+~~~~~~~~~~~~~~~~~
+
+**TensorDict-level** (these control lines in the outer ``TensorDict(...)``
+block):
+
+====================  ===========  ===================================================
+Option                Default      Description
+====================  ===========  ===================================================
+``show_batch_size``   ``True``     Show the ``batch_size=`` line.
+``show_device``       ``True``     Show the ``device=`` line.
+``show_is_shared``    ``True``     Show the ``is_shared=`` line.
+====================  ===========  ===================================================
+
+**Tensor-level** (these control what appears inside each
+``Tensor(...)`` field descriptor):
+
+==========================  ===========  ============================================
+Option                      Default      Description
+==========================  ===========  ============================================
+``show_shape``              ``True``     Show the ``shape=`` attribute.
+``show_tensor_device``      ``True``     Show the ``device=`` attribute.
+``show_tensor_dtype``       ``True``     Show the ``dtype=`` attribute.
+``show_tensor_is_shared``   ``True``     Show the ``is_shared=`` attribute.
+==========================  ===========  ============================================
+
+**Extended attributes** (off by default -- opt-in for deeper debugging):
+
+======================  ===========  =================================================
+Option                  Default      Description
+======================  ===========  =================================================
+``show_grad``           ``False``    Show ``requires_grad=``.
+``show_is_contiguous``  ``False``    Show ``is_contiguous=``.
+``show_is_view``        ``False``    Show ``is_view=`` (whether ``._base`` is set).
+``show_storage_size``   ``False``    Show ``storage_size=`` (bytes).
+``plain``               ``False``    Append a short value summary
+                                     (mean/std for floats, min/max for ints).
+======================  ===========  =================================================
+
+Querying the current settings
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+:func:`~tensordict.get_printoptions` returns a dict with the current values:
+
+    >>> from tensordict import get_printoptions
+    >>> get_printoptions()
+    {'show_batch_size': True, 'show_device': True, ...}
+
+
+Reconstructing a TensorDict from its printed representation
+-----------------------------------------------------------
+
+When debugging, it is common to receive a TensorDict repr as a string -- for
+example, pasted from a log file or a colleague's terminal.
+:func:`~tensordict.parse_tensor_dict_string` can reconstruct a dummy
+:class:`~tensordict.TensorDict` from that string.  The resulting object has the
+correct structure, batch size, device and dtypes, but all tensor values are
+replaced by zeros (since the repr does not contain actual data):
+
+    >>> from tensordict import parse_tensor_dict_string
+    >>> s = """TensorDict(
+    ...     fields={
+    ...         image: Tensor(shape=torch.Size([32, 3, 64, 64]), device=cpu, dtype=torch.float32, is_shared=False),
+    ...         label: Tensor(shape=torch.Size([32]), device=cpu, dtype=torch.int64, is_shared=False)},
+    ...     batch_size=torch.Size([32]),
+    ...     device=cpu,
+    ...     is_shared=False)"""
+    >>> td = parse_tensor_dict_string(s)
+    >>> td.batch_size
+    torch.Size([32])
+    >>> td["image"].shape
+    torch.Size([32, 3, 64, 64])
+
+.. note::
+
+    :func:`~tensordict.parse_tensor_dict_string` currently only works with the
+    default (``plain``) print format -- the one that includes ``shape=``,
+    ``device=``, ``dtype=`` and ``is_shared=`` for every field.
+    If attributes have been hidden via :class:`~tensordict.set_printoptions`,
+    the regex parser will not find the expected fields and reconstruction will
+    fail.  Support for non-default formats will be added in a follow-up PR.

docs/source/reference/td.rst

@@ -288,4 +288,6 @@ Utils
     set_capture_non_tensor_stack
     set_lazy_legacy
     set_list_to_stack
+    set_printoptions
+    get_printoptions
     list_to_stack

tensordict/_lazy.py

@@ -3515,9 +3515,7 @@ def __repr__(self):
         fields = _td_fields(self)
         parts = [
             indent(f"fields={{{fields}}}", 4 * " "),
-            indent(
-                f"exclusive_fields={{{self._repr_exclusive_fields()}}}", 4 * " "
-            ),
+            indent(f"exclusive_fields={{{self._repr_exclusive_fields()}}}", 4 * " "),
         ]
         if _REPR_OPTIONS["show_batch_size"]:
             parts.append(indent(f"batch_size={self.batch_size}", 4 * " "))

tensordict/base.py

@@ -81,8 +81,8 @@
     _prefix_last_key,
     _proc_init,
     _prune_selected_keys,
-    _REPR_OPTIONS,
     _rebuild_njt_from_njt,
+    _REPR_OPTIONS,
     _set_max_batch_size,
     _shape,
     _split_tensordict,

test/test_tensordict.py

@@ -10194,7 +10194,12 @@ def test_hide_multiple(self):
         from tensordict import set_printoptions
 
         td = TensorDict({"a": torch.randn(3, 4)})
-        with set_printoptions(show_device=False, show_is_shared=False, show_tensor_dtype=False, show_tensor_is_shared=False):
+        with set_printoptions(
+            show_device=False,
+            show_is_shared=False,
+            show_tensor_dtype=False,
+            show_tensor_is_shared=False,
+        ):
             r = repr(td)
         assert "\n    device=" not in r
         assert "\n    is_shared=" not in r
@@ -10286,7 +10291,12 @@ class MyClass:
             x: torch.Tensor
 
         obj = MyClass(x=torch.randn(3, 4), batch_size=[3])
-        with set_printoptions(show_device=False, show_is_shared=False, show_tensor_device=False, show_tensor_is_shared=False):
+        with set_printoptions(
+            show_device=False,
+            show_is_shared=False,
+            show_tensor_device=False,
+            show_tensor_is_shared=False,
+        ):
             r = repr(obj)
         assert "MyClass(" in r
         assert "device=" not in r