Add env var option to override torch.Tensor.__repr__ (#384)

The __repr__ for large tensors is way too slow and the debugger takes
forever to access them stalling for minutes at a time. This gives the
option to replace it.

sharktank/sharktank/types/tensors.py

@@ -19,6 +19,7 @@
 from copy import deepcopy
 from collections.abc import Collection, Sequence
 from numbers import Integral, Number
+import os
 
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
@@ -55,6 +56,18 @@
     "UnreducedTensor",
 ]
 
+if (
+    "SHARKTANK_OVERRIDE_TORCH_TENSOR_REPR" in os.environ
+    and os.environ["SHARKTANK_OVERRIDE_TORCH_TENSOR_REPR"] != "0"
+):
+
+    def _tensor_debugger_friendly_repr(self: torch.Tensor):
+        """Override for the torch.Tensor.__repr__ so it does not take forever when the
+        debugger wants to query many/large tensors."""
+        return f"Tensor({list(self.shape)}, {self.dtype})"
+
+    Tensor.__repr__ = _tensor_debugger_friendly_repr
+
 # JSON encodable value types.
 MetaDataValueType = Union[int, bool, float, str]
 UnnamedTensorName = "<unnamed>"