Merge pull request #21 from mctigger/tree_map-error-messages

Improved PyTree Error Messages and Refactoring

Author: mctigger

examples/tensor_dataclass/04_stack.py

@@ -25,6 +25,10 @@ class DataPoint(TensorDataClass):
     y: torch.Tensor
 
 
+class NotADataPoint(TensorDataClass):
+    a: torch.Tensor
+
+
 def main() -> None:
     """Demonstrate stacking TensorDataClass instances."""
     point1 = DataPoint(
@@ -49,6 +53,13 @@ def main() -> None:
     assert stacked.x.shape == (2, 3, 4)
     assert stacked.y.shape == (2, 3, 5)
 
+    # Attempt to stack with different TensorDataClasses
+    try:
+        not_a_point = NotADataPoint(a=torch.rand(3, 4), shape=(3,), device="cpu")
+        torch.stack([point1, not_a_point], dim=0)
+    except Exception as e:
+        print(e)
+
 
 if __name__ == "__main__":
     main()

src/tensorcontainer/tensor_annotated.py

@@ -1,12 +1,16 @@
 from __future__ import annotations
 
+from dataclasses import dataclass
 from typing import Any, Iterable, TypeVar, Union, get_args
 
 from torch import Tensor
 from torch.utils import _pytree as pytree
 from typing_extensions import Self
 
-from tensorcontainer.tensor_container import TensorContainer
+from tensorcontainer.tensor_container import (
+    TensorContainer,
+    TensorContainerPytreeContext,
+)
 from tensorcontainer.types import DeviceLike, ShapeLike
 from tensorcontainer.utils import PytreeRegistered
 
@@ -17,6 +21,54 @@
 T_TensorAnnotated = TypeVar("T_TensorAnnotated", bound="TensorAnnotated")
 
 
+# PyTree context metadata for reconstruction
+@dataclass
+class TensorAnnotatedPytreeContext(
+    TensorContainerPytreeContext["TensorAnnotatedPytreeContext"]
+):
+    """TensorAnnotated PyTree context with enhanced error messages."""
+
+    keys: list[str]
+    event_ndims: list[int]
+    metadata: dict[str, Any]
+
+    def __str__(self) -> str:
+        """Return human-readable description of this TensorAnnotated context."""
+        # Try to get the actual class name from metadata
+        class_name = self.metadata.get("class_name", "TensorDataClass")
+
+        fields_str = f"fields={self.keys}"
+        device_str = f"device={self.device}"
+
+        return f"{class_name}({fields_str}, {device_str})"
+
+    def analyze_mismatch_with(
+        self, other: TensorAnnotatedPytreeContext, entry_index: int
+    ) -> str:
+        """Analyze specific mismatches between TensorAnnotated contexts."""
+        # Start with base class analysis (device mismatch, if any)
+        guidance = super().analyze_mismatch_with(other, entry_index)
+
+        # Add TensorAnnotated-specific analysis
+        self_fields = set(self.keys)
+        other_fields = set(other.keys)
+
+        if self_fields != other_fields:
+            missing = self_fields - other_fields
+            extra = other_fields - self_fields
+            guidance += "Field mismatch detected."
+            if missing:
+                guidance += (
+                    f" Missing fields in container {entry_index}: {sorted(missing)}."
+                )
+            if extra:
+                guidance += (
+                    f" Extra fields in container {entry_index}: {sorted(extra)}."
+                )
+
+        return guidance
+
+
 class TensorAnnotated(TensorContainer, PytreeRegistered):
     def __init__(self, shape: ShapeLike, device: DeviceLike | None):
         super().__init__(shape, device)
@@ -73,11 +125,13 @@ def _get_meta_attributes(self):
 
     def _get_pytree_context(
         self, flat_names: list[str], flat_leaves: list[TDCompatible], meta_data
-    ) -> tuple:
+    ) -> TensorAnnotatedPytreeContext:
         batch_ndim = len(self.shape)
-        event_ndims = tuple(leaf.ndim - batch_ndim for leaf in flat_leaves)
+        event_ndims = [leaf.ndim - batch_ndim for leaf in flat_leaves]
 
-        return flat_names, event_ndims, meta_data, self.device
+        return TensorAnnotatedPytreeContext(
+            self.device, flat_names, event_ndims, meta_data
+        )
 
     def _pytree_flatten(self) -> tuple[list[Any], Any]:
         tensor_attributes = self._get_tensor_attributes()
@@ -94,15 +148,20 @@ def _pytree_flatten_with_keys_fn(
         self,
     ) -> tuple[list[tuple[pytree.KeyEntry, Any]], Any]:
         flat_values, context = self._pytree_flatten()
-        flat_names = context[0]
+        flat_names = context.keys
         name_value_tuples = [
             (pytree.GetAttrKey(k), v) for k, v in zip(flat_names, flat_values)
         ]
         return name_value_tuples, context  # type: ignore[return-value]
 
     @classmethod
-    def _pytree_unflatten(cls, leaves: Iterable[Any], context: pytree.Context) -> Self:
-        flat_names, event_ndims, meta_data, device = context
+    def _pytree_unflatten(
+        cls, leaves: Iterable[Any], context: TensorAnnotatedPytreeContext
+    ) -> Self:
+        flat_names = context.keys
+        event_ndims = context.event_ndims
+        device = context.device
+        meta_data = context.metadata
 
         leaves = list(leaves)  # Convert to list to allow indexing
 

src/tensorcontainer/tensor_container.py

@@ -1,11 +1,23 @@
 from __future__ import annotations
 
+from dataclasses import dataclass
 import functools
 import textwrap
 import threading
-from abc import abstractmethod
+from abc import ABC, abstractmethod
 from contextlib import contextmanager
-from typing import Any, Callable, Iterable, List, Optional, Tuple, Type, Union
+from typing import (
+    Any,
+    Callable,
+    Generic,
+    Iterable,
+    List,
+    Optional,
+    Tuple,
+    Type,
+    TypeVar,
+    Union,
+)
 
 import torch
 
@@ -15,7 +27,12 @@
 from typing_extensions import Self, TypeAlias
 
 from tensorcontainer.types import DeviceLike, ShapeLike
-from tensorcontainer.utils import resolve_device
+from tensorcontainer.utils import (
+    ContextWithAnalysis,
+    diagnose_pytree_structure_mismatch,
+    resolve_device,
+    format_path,
+)
 
 HANDLED_FUNCTIONS = {}
 
@@ -33,6 +50,25 @@ def decorator(func):
     return decorator
 
 
+U = TypeVar("U", bound="TensorContainerPytreeContext")
+
+
+@dataclass
+class TensorContainerPytreeContext(ContextWithAnalysis[U], Generic[U], ABC):
+    """Base PyTree context class for tensor containers with common device handling."""
+
+    device: torch.device | None
+
+    def analyze_mismatch_with(self, other: U, entry_index: int) -> str:
+        """Analyze mismatches with another TensorContainerPytreeContext, starting with device analysis."""
+        # Check device mismatch first
+        if self.device != other.device:
+            return f"Device mismatch: container 0 device={self.device}, container {entry_index} device={other.device}. "
+
+        # If devices match, return empty string for subclasses to add their analysis
+        return ""
+
+
 class TensorContainer:
     """A foundational base class for PyTree-compatible tensor containers with batch semantics.
 
@@ -308,11 +344,25 @@ def wrapped_func(keypath, x, *xs):
             try:
                 return func(x, *xs)
             except Exception as e:
-                path = cls._format_path(keypath)
+                path = format_path(keypath)
                 message = f"Error at path {path}: {type(e).__name__}: {e}"
                 raise type(e)(message) from e
 
-        return cls.tree_map_with_path(wrapped_func, tree, *rests, is_leaf=is_leaf)
+        try:
+            return pytree.tree_map_with_path(
+                wrapped_func, tree, *rests, is_leaf=is_leaf
+            )
+        except Exception as e:
+            # The following code is just to provide better error messages for operations that
+            # work on multiple pytrees such as torch.stack() or torch.cat()
+            # It is not necessary for TensorContainer to function properly.
+            if len(rests) > 0:
+                msg = diagnose_pytree_structure_mismatch(tree, *rests, is_leaf=is_leaf)
+                if msg:
+                    raise RuntimeError(msg) from e
+
+            # Re-raise if it is an unknown error.
+            raise e
 
     @classmethod
     def tree_map_with_path(
@@ -431,22 +481,6 @@ def transform_ellipsis_index(self, shape: torch.Size, idx: tuple) -> tuple:
 
         return final_index
 
-    @classmethod
-    def _format_path(cls, path: pytree.KeyPath) -> str:
-        """Helper to format a PyTree KeyPath into a readable string."""
-        parts = []
-        for entry in path:
-            if isinstance(entry, tuple):  # Handle nested KeyPath tuples
-                parts.append(cls._format_path(entry))
-            else:
-                parts.append(str(entry))
-
-        # Join parts and clean up leading dots if any
-        formatted_path = "".join(parts)
-        if formatted_path.startswith("."):
-            formatted_path = formatted_path[1:]
-        return formatted_path
-
     def __repr__(self) -> str:
         # Use a consistent indent of 4 spaces, which is standard
         indent = "    "
@@ -863,14 +897,10 @@ def _stack(
         dim = dim + batch_ndim + 1
 
     if dim < 0 or dim > batch_ndim:
-        raise IndexError("Dimension out of range")
-
-    shape_expected = first_tc.shape
-
-    for t in tensors:
-        shape_is = t.shape
-        if shape_is != shape_expected:
-            raise ValueError("stack expects each TensorContainer to be equal size")
+        raise IndexError(
+            f"Dimension {dim - batch_ndim - 1 if dim < 0 else dim} out of range "
+            f"(expected 0 to {batch_ndim} for stack operation on shape {tuple(first_tc.shape)})"
+        )
 
     # Pytree handles the stacking of individual tensors and metadata consistency
     result_td = TensorContainer._tree_map(lambda *x: torch.stack(x, dim), *tensors)
@@ -891,16 +921,10 @@ def _cat(
         dim = dim + batch_ndim
 
     if dim < 0 or dim > batch_ndim - 1:
-        raise IndexError("Dimension out of range")
-
-    shape_expected = first_tc.shape[:dim] + first_tc.shape[dim + 1 :]
-
-    for t in tensors:
-        shape_is = t.shape[:dim] + t.shape[dim + 1 :]
-        if shape_is != shape_expected:
-            raise ValueError(
-                "TensorContainer batch shapes must be identical except for 'dim'"
-            )
+        raise IndexError(
+            f"Dimension {dim - batch_ndim if dim < 0 else dim} out of range "
+            f"(expected 0 to {batch_ndim - 1} for concatenation on shape {tuple(first_tc.shape)})"
+        )
 
     # Create a new TensorContainer of the same type as the first one
     # and apply torch.cat to its internal tensors

src/tensorcontainer/tensor_dict.py

@@ -16,13 +16,13 @@
 
 from __future__ import annotations
 
+from dataclasses import dataclass
 from typing import (
     Any,
     Dict,
     Iterable,
     List,
     Mapping,
-    NamedTuple,
     Tuple,
     Union,
     cast,
@@ -38,20 +38,56 @@
     PyTree,
 )
 
-from tensorcontainer.tensor_container import TensorContainer
+from tensorcontainer.tensor_container import (
+    TensorContainer,
+    TensorContainerPytreeContext,
+)
 from tensorcontainer.types import DeviceLike, ShapeLike
 from tensorcontainer.utils import PytreeRegistered
 
 TDCompatible = Union[Tensor, TensorContainer]
 
 
 # PyTree context metadata for reconstruction
-class TensorDictPytreeContext(NamedTuple):
-    keys: Tuple[str, ...]
-    event_ndims: Tuple[int, ...]
-    device_context: torch.device | None
+@dataclass
+class TensorDictPytreeContext(TensorContainerPytreeContext["TensorDictPytreeContext"]):
+    """TensorDict PyTree context with enhanced error messages."""
+
+    keys: list[str]
+    event_ndims: list[int]
     metadata: Dict[str, Any]
 
+    def __str__(self) -> str:
+        """Return human-readable description of this TensorDict context."""
+        keys_str = f"keys={list(self.keys)}" if self.keys else "keys=[]"
+        device_str = f"device={self.device}" if self.device else "device=None"
+
+        return f"TensorDict({keys_str}, {device_str})"
+
+    def analyze_mismatch_with(
+        self, other: TensorDictPytreeContext, entry_index: int
+    ) -> str:
+        """Analyze specific mismatches with another TensorDict context."""
+        # Start with base class analysis (device mismatch, if any)
+        guidance = super().analyze_mismatch_with(other, entry_index)
+
+        # Add TensorDict-specific analysis
+        self_keys = set(self.keys)
+        other_keys = set(other.keys)
+
+        if self_keys != other_keys:
+            missing = self_keys - other_keys
+            extra = other_keys - self_keys
+            guidance += "Key mismatch detected."
+            if missing:
+                guidance += (
+                    f" Missing keys in container {entry_index}: {sorted(missing)}."
+                )
+            if extra:
+                guidance += f" Extra keys in container {entry_index}: {sorted(extra)}."
+
+        return guidance
+
 
 class TensorDict(TensorContainer, PytreeRegistered):
     """Dictionary-like container for batched tensors that share the same leading batch shape.
@@ -159,7 +195,7 @@ def _get_pytree_context(
         """
         batch_ndim = len(self.shape)
         event_ndims = tuple(leaf.ndim - batch_ndim for leaf in flat_leaves)
-        return TensorDictPytreeContext(tuple(keys), event_ndims, self.device, metadata)
+        return TensorDictPytreeContext(self.device, tuple(keys), event_ndims, metadata)
 
     def _pytree_flatten(
         self,
@@ -225,8 +261,11 @@ def _pytree_unflatten(
           - If no leaves are provided, an empty ``TensorDict`` is constructed using the shape
             from the context. The device is restored from the context.
         """
-        # Unpack context tuple
-        keys, event_ndims, device_context, metadata = context
+        # Access context fields
+        keys = context.keys
+        event_ndims = context.event_ndims
+        device_context = context.device
+        metadata = context.metadata
 
         obj = cls.__new__(cls)
         obj.device = device_context