init

pytorch · Jan 13, 2025 · 4713734 · 4713734
1 parent c750598
commit 4713734
Showing 1 changed file with 44 additions and 8 deletions.
diff --git a/tensordict/return_types.py b/tensordict/return_types.py
@@ -7,33 +7,69 @@
 from tensordict.tensordict import TensorDict
 
 
-@tensorclass
+@tensorclass(shadow=True)
 class min:
     """A `min` tensorclass to be used as a result for :meth:`~tensordict.TensorDict.min` operations."""
 
-    vals: TensorDict
+    values: TensorDict
     indices: TensorDict
 
+    def __getitem__(self, item):
+        try:
+            return (self.values, self.indices)[item]
+        except IndexError:
+            raise IndexError(
+                f"Indexing a {type(self)} element follows the torch.return_types.{type(self).__name__}'s "
+                f"__getitem__ method API."
+            )
 
-@tensorclass
+
+@tensorclass(shadow=True)
 class max:
     """A `max` tensorclass to be used as a result for :meth:`~tensordict.TensorDict.max` operations."""
 
-    vals: TensorDict
+    values: TensorDict
     indices: TensorDict
 
+    def __getitem__(self, item):
+        try:
+            return (self.values, self.indices)[item]
+        except IndexError:
+            raise IndexError(
+                f"Indexing a {type(self)} element follows the torch.return_types.{type(self).__name__}'s "
+                f"__getitem__ method API."
+            )
+
 
-@tensorclass
+@tensorclass(shadow=True)
 class cummin:
     """A `cummin` tensorclass to be used as a result for :meth:`~tensordict.TensorDict.cummin` operations."""
 
-    vals: TensorDict
+    values: TensorDict
     indices: TensorDict
 
+    def __getitem__(self, item):
+        try:
+            return (self.values, self.indices)[item]
+        except IndexError:
+            raise IndexError(
+                f"Indexing a {type(self)} element follows the torch.return_types.{type(self).__name__}'s "
+                f"__getitem__ method API."
+            )
 
-@tensorclass
+
+@tensorclass(shadow=True)
 class cummax:
     """A `cummax` tensorclass to be used as a result for :meth:`~tensordict.TensorDict.cummax` operations."""
 
-    vals: TensorDict
+    values: TensorDict
     indices: TensorDict
+
+    def __getitem__(self, item):
+        try:
+            return (self.values, self.indices)[item]
+        except IndexError:
+            raise IndexError(
+                f"Indexing a {type(self)} element follows the torch.return_types.{type(self).__name__}'s "
+                f"__getitem__ method API."
+            )