init

Author: vmoens

tensordict/__init__.py

@@ -60,6 +60,7 @@
     assert_allclose_td,
     assert_close,
     capture_non_tensor_stack,
+    get_printoptions,
     is_batchedtensor,
     is_non_tensor,
     is_tensorclass,
@@ -69,6 +70,7 @@
     set_capture_non_tensor_stack,
     set_lazy_legacy,
     set_list_to_stack,
+    set_printoptions,
     unravel_key,
     unravel_key_list,
 )
@@ -159,6 +161,8 @@
     "set_lazy_legacy",
     "list_to_stack",
     "set_list_to_stack",
+    "set_printoptions",
+    "get_printoptions",
     # TensorClass components
     "tensorclass",
     "MetaData",

tensordict/_lazy.py

@@ -66,6 +66,7 @@
     _parse_to,
     _recursive_unbind_list,
     _renamed_inplace_method,
+    _REPR_OPTIONS,
     _shape,
     _td_fields,
     _unravel_key_to_tuple,
@@ -3512,24 +3513,20 @@ def _propagate_unlock(self):
 
     def __repr__(self):
         fields = _td_fields(self)
-        field_str = indent(f"fields={{{fields}}}", 4 * " ")
-        exclusive_fields_str = indent(
-            f"exclusive_fields={{{self._repr_exclusive_fields()}}}", 4 * " "
-        )
-        batch_size_str = indent(f"batch_size={self.batch_size}", 4 * " ")
-        device_str = indent(f"device={self.device}", 4 * " ")
-        is_shared_str = indent(f"is_shared={self.is_shared()}", 4 * " ")
-        stack_dim = indent(f"stack_dim={self.stack_dim}", 4 * " ")
-        string = ",\n".join(
-            [
-                field_str,
-                exclusive_fields_str,
-                batch_size_str,
-                device_str,
-                is_shared_str,
-                stack_dim,
-            ]
-        )
+        parts = [
+            indent(f"fields={{{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 * " "))
+        if _REPR_OPTIONS["show_device"]:
+            parts.append(indent(f"device={self.device}", 4 * " "))
+        if _REPR_OPTIONS["show_is_shared"]:
+            parts.append(indent(f"is_shared={self.is_shared()}", 4 * " "))
+        parts.append(indent(f"stack_dim={self.stack_dim}", 4 * " "))
+        string = ",\n".join(parts)
         return f"{type(self).__name__}(\n{string})"
 
     def _exclusive_keys(self):

tensordict/base.py

@@ -81,6 +81,7 @@
     _prefix_last_key,
     _proc_init,
     _prune_selected_keys,
+    _REPR_OPTIONS,
     _rebuild_njt_from_njt,
     _set_max_batch_size,
     _shape,
@@ -553,11 +554,14 @@ def __lt__(self, other: object) -> Self:
     def __repr__(self) -> str:
         try:
             fields = _td_fields(self)
-            field_str = indent(f"fields={{{fields}}}", 4 * " ")
-            batch_size_str = indent(f"batch_size={self.batch_size}", 4 * " ")
-            device_str = indent(f"device={self.device}", 4 * " ")
-            is_shared_str = indent(f"is_shared={self.is_shared()}", 4 * " ")
-            string = ",\n".join([field_str, batch_size_str, device_str, is_shared_str])
+            parts = [indent(f"fields={{{fields}}}", 4 * " ")]
+            if _REPR_OPTIONS["show_batch_size"]:
+                parts.append(indent(f"batch_size={self.batch_size}", 4 * " "))
+            if _REPR_OPTIONS["show_device"]:
+                parts.append(indent(f"device={self.device}", 4 * " "))
+            if _REPR_OPTIONS["show_is_shared"]:
+                parts.append(indent(f"is_shared={self.is_shared()}", 4 * " "))
+            string = ",\n".join(parts)
         except AttributeError:
             # When using torch.compile, an exception may be raised with a tensordict object
             #  that has no attribute (no _tensordict or no _batch_size).

tensordict/tensorclass.py

@@ -63,6 +63,7 @@
     _is_json_serializable,
     _is_tensorclass,
     _LOCK_ERROR,
+    _REPR_OPTIONS,
     _td_fields,
     _TENSORCLASS_MEMO,
     _unravel_key_to_tuple,
@@ -2274,34 +2275,31 @@ def _repr(self) -> str:
     field_str = [fields] if fields else []
     non_tensor_fields = _all_non_td_fields_as_str(self._non_tensordict)
 
-    medatada_fields = []
+    metadata_fields = []
 
-    if "batch_size" not in self.__expected_keys__:
-        batch_size_str = indent(f"batch_size={self.batch_size}", 4 * " ")
-        medatada_fields.append(batch_size_str)
-    elif "shape" not in self.__expected_keys__:
-        batch_size_str = indent(f"shape={self.shape}", 4 * " ")
-        medatada_fields.append(batch_size_str)
-    if "device" not in self.__expected_keys__:
-        device_str = indent(f"device={self.device}", 4 * " ")
-        medatada_fields.append(device_str)
-
-    is_shared_str = indent(f"is_shared={self.is_shared()}", 4 * " ")
-    medatada_fields.append(is_shared_str)
+    if _REPR_OPTIONS["show_batch_size"]:
+        if "batch_size" not in self.__expected_keys__:
+            metadata_fields.append(indent(f"batch_size={self.batch_size}", 4 * " "))
+        elif "shape" not in self.__expected_keys__:
+            metadata_fields.append(indent(f"shape={self.shape}", 4 * " "))
+    if _REPR_OPTIONS["show_device"] and "device" not in self.__expected_keys__:
+        metadata_fields.append(indent(f"device={self.device}", 4 * " "))
+    if _REPR_OPTIONS["show_is_shared"]:
+        metadata_fields.append(indent(f"is_shared={self.is_shared()}", 4 * " "))
 
     if len(non_tensor_fields) > 0:
         non_tensor_field_str = indent(
             ",\n".join(non_tensor_fields),
             4 * " ",
         )
         if field_str:
-            string = ",\n".join(field_str + [non_tensor_field_str, *medatada_fields])
+            string = ",\n".join(field_str + [non_tensor_field_str, *metadata_fields])
         else:
-            string = ",\n".join([non_tensor_field_str, *medatada_fields])
+            string = ",\n".join([non_tensor_field_str, *metadata_fields])
     elif field_str:
-        string = ",\n".join(field_str + medatada_fields)
-    elif len(medatada_fields) > 0:
-        string = ",\n".join(medatada_fields)
+        string = ",\n".join(field_str + metadata_fields)
+    elif len(metadata_fields) > 0:
+        string = ",\n".join(metadata_fields)
     else:
         string = ""
     return f"{type(self).__name__}({string})"

tensordict/utils.py

@@ -1361,27 +1361,62 @@ def istensor(cls):
     return True
 
 
+def _plain_summary(tensor: Tensor) -> str:
+    """Short value summary for ``plain`` print mode."""
+    if isinstance(tensor, UninitializedTensorMixin) or tensor.numel() == 0:
+        return "[]"
+    if tensor.is_floating_point() or tensor.is_complex():
+        return f"mean={tensor.to(torch.float64).mean().item():.4g}, std={tensor.to(torch.float64).std().item():.4g}"
+    return f"min={tensor.min().item()}, max={tensor.max().item()}"
+
+
+def _tensor_repr_fields(tensor: Tensor) -> list[str]:
+    """Build the list of ``key=value`` strings for a single tensor descriptor."""
+    opts = _REPR_OPTIONS
+    parts: list[str] = []
+    if opts["show_shape"]:
+        parts.append(f"shape={_shape(tensor)}")
+    if opts["show_tensor_device"]:
+        parts.append(f"device={_device(tensor)}")
+    if opts["show_tensor_dtype"]:
+        parts.append(f"dtype={_dtype(tensor)}")
+    if opts["show_tensor_is_shared"]:
+        parts.append(f"is_shared={_is_shared(tensor)}")
+    if opts["show_grad"]:
+        parts.append(f"requires_grad={tensor.requires_grad}")
+    if opts["show_is_contiguous"]:
+        parts.append(f"is_contiguous={tensor.is_contiguous()}")
+    if opts["show_is_view"]:
+        parts.append(f"is_view={tensor._base is not None}")
+    if opts["show_storage_size"]:
+        parts.append(f"storage_size={tensor.untyped_storage().nbytes()}")
+    if opts["plain"]:
+        parts.append(_plain_summary(tensor))
+    return parts
+
+
 def _get_repr(tensor: Tensor) -> str:
-    s = ", ".join(
-        [
-            f"shape={_shape(tensor)}",
-            f"device={_device(tensor)}",
-            f"dtype={_dtype(tensor)}",
-            f"is_shared={_is_shared(tensor)}",
-        ]
-    )
+    s = ", ".join(_tensor_repr_fields(tensor))
     return f"{type(tensor).__name__}({s})"
 
 
+def _tensor_repr_fields_custom(shape, device, dtype, is_shared) -> list[str]:
+    """Build the list of ``key=value`` strings for a tensor descriptor with explicit values."""
+    opts = _REPR_OPTIONS
+    parts: list[str] = []
+    if opts["show_shape"]:
+        parts.append(f"shape={shape}")
+    if opts["show_tensor_device"]:
+        parts.append(f"device={device}")
+    if opts["show_tensor_dtype"]:
+        parts.append(f"dtype={dtype}")
+    if opts["show_tensor_is_shared"]:
+        parts.append(f"is_shared={is_shared}")
+    return parts
+
+
 def _get_repr_custom(cls, shape, device, dtype, is_shared) -> str:
-    s = ", ".join(
-        [
-            f"shape={shape}",
-            f"device={device}",
-            f"dtype={dtype}",
-            f"is_shared={is_shared}",
-        ]
-    )
+    s = ", ".join(_tensor_repr_fields_custom(shape, device, dtype, is_shared))
     return f"{cls.__name__}({s})"
 
 
@@ -1766,6 +1801,107 @@ def _getitem_batch_size(batch_size, index):
     return torch.Size(out)
 
 
+# Repr / print options
+_REPR_OPTIONS = {
+    "show_batch_size": True,
+    "show_device": True,
+    "show_is_shared": True,
+    "show_shape": True,
+    "show_tensor_device": True,
+    "show_tensor_dtype": True,
+    "show_tensor_is_shared": True,
+    "show_grad": False,
+    "show_is_contiguous": False,
+    "show_is_view": False,
+    "show_storage_size": False,
+    "plain": False,
+}
+
+_REPR_OPTIONS_KEYS = frozenset(_REPR_OPTIONS)
+
+
+class set_printoptions(_DecoratorContextManager):
+    """Controls which attributes appear in TensorDict's ``__repr__`` output.
+
+    Can be used as a global setter (via :meth:`set`), a context manager, or a
+    decorator.  Follows the same pattern as :class:`set_lazy_legacy`.
+
+    Keyword Args:
+        show_batch_size (bool, optional): Show ``batch_size`` in TensorDict repr.
+            Defaults to ``True``.
+        show_device (bool, optional): Show ``device`` in TensorDict repr.
+            Defaults to ``True``.
+        show_is_shared (bool, optional): Show ``is_shared`` in TensorDict repr.
+            Defaults to ``True``.
+        show_shape (bool, optional): Show ``shape`` in per-tensor field descriptors.
+            Defaults to ``True``.
+        show_tensor_device (bool, optional): Show ``device`` in per-tensor field
+            descriptors. Defaults to ``True``.
+        show_tensor_dtype (bool, optional): Show ``dtype`` in per-tensor field
+            descriptors. Defaults to ``True``.
+        show_tensor_is_shared (bool, optional): Show ``is_shared`` in per-tensor
+            field descriptors. Defaults to ``True``.
+        show_grad (bool, optional): Show ``requires_grad`` in per-tensor field
+            descriptors. Defaults to ``False``.
+        show_is_contiguous (bool, optional): Show ``is_contiguous`` in per-tensor
+            field descriptors. Defaults to ``False``.
+        show_is_view (bool, optional): Show ``is_view`` in per-tensor field
+            descriptors. Defaults to ``False``.
+        show_storage_size (bool, optional): Show ``storage_size`` (in bytes) in
+            per-tensor field descriptors. Defaults to ``False``.
+        plain (bool, optional): When ``True``, include a short summary of the
+            actual tensor values in the field descriptors. Defaults to ``False``.
+
+    Examples:
+        >>> import torch
+        >>> from tensordict import TensorDict, set_printoptions
+        >>> td = TensorDict({"x": torch.randn(3, 4)})
+        >>> # Global
+        >>> set_printoptions(show_device=False, show_is_shared=False).set()
+        >>> print(td)
+        >>> # Context manager
+        >>> with set_printoptions(show_tensor_dtype=False):
+        ...     print(td)
+        >>> # Decorator
+        >>> @set_printoptions(show_is_shared=False)
+        ... def my_func(td):
+        ...     print(td)
+
+    """
+
+    def __init__(self, **kwargs: bool) -> None:
+        super().__init__()
+        unknown = set(kwargs) - _REPR_OPTIONS_KEYS
+        if unknown:
+            raise TypeError(
+                f"Unknown printoptions: {unknown}. Valid options: {sorted(_REPR_OPTIONS_KEYS)}"
+            )
+        self._kwargs = kwargs
+
+    def clone(self) -> set_printoptions:
+        return type(self)(**self._kwargs)
+
+    def __enter__(self) -> None:
+        self.set()
+
+    def set(self) -> None:
+        global _REPR_OPTIONS
+        self._old = dict(_REPR_OPTIONS)
+        _REPR_OPTIONS.update(self._kwargs)
+
+    def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
+        global _REPR_OPTIONS
+        _REPR_OPTIONS.update(self._old)
+
+
+def get_printoptions() -> dict:
+    """Returns the current TensorDict print options as a dict.
+
+    See :class:`set_printoptions` for details on each option.
+    """
+    return dict(_REPR_OPTIONS)
+
+
 # Lazy classes control (legacy feature)
 _DEFAULT_LAZY_OP = False
 _LAZY_OP = os.environ.get("LAZY_LEGACY_OP")