diff --git a/tools/pyi/gen_pyi.py b/tools/pyi/gen_pyi.py index 4154188f058..24f6b46549c 100644 --- a/tools/pyi/gen_pyi.py +++ b/tools/pyi/gen_pyi.py @@ -433,6 +433,12 @@ def gen_pyi( "_is_functional_tensor": [ "def _is_functional_tensor(t: Tensor) -> _bool: ..." ], + "_from_functional_tensor": [ + "def _from_functional_tensor(t: Tensor) -> Tensor: ..." + ], + "_to_functional_tensor": [ + "def _to_functional_tensor(t: Tensor) -> Tensor: ..." + ], "range": [ "def range(start: Number, end: Number," " step: Number=1, *, out: Optional[Tensor]=None, {}) -> Tensor: ...".format( diff --git a/torch/_tensor_str.py b/torch/_tensor_str.py index a9b83f31b93..54758a530af 100644 --- a/torch/_tensor_str.py +++ b/torch/_tensor_str.py @@ -415,6 +415,9 @@ def _str_intern(inp, *, tensor_contents=None): return "\n".join(f" {line}" for line in s.split("\n")) strs = ",\n".join(indented_str(str(t), indent + 1) for t in torch.ops.aten.unbind.int(self, 0)) tensor_str = f"[\n{strs}\n]" + elif torch._is_functional_tensor(self): + prefix = "_to_functional_tensor(" + tensor_str = repr(torch._from_functional_tensor(self)) else: if self.is_meta: suffixes.append('size=' + str(tuple(self.shape)))