OSSForks/pytorch
1
0
Fork 0
mirror of https://github.com/zebrajr/pytorch.git synced 2025-12-06 12:20:52 +01:00
Code Issues Actions 38 Packages Projects Releases Wiki Activity

Make functional tensors printable (#81454)

Browse Source 
Signed-off-by: Edward Z. Yang <ezyang@fb.com>
Pull Request resolved: https://github.com/pytorch/pytorch/pull/81454
Approved by: https://github.com/bdhirsh
This commit is contained in:
Edward Z. Yang 2022-07-13 23:17:23 -04:00 committed by PyTorch MergeBot
parent 6f6e47aba4
commit 62f1ff23fb
2 changed files with 9 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
tools/pyi/gen_pyi.py
View File

@ -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(

3
torch/_tensor_str.py
View File

@ -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)))