mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-21 05:34:18 +08:00
7c12cc7ce4
See #127836 for details. Pull Request resolved: https://github.com/pytorch/pytorch/pull/127843 Approved by: https://github.com/oulgen ghstack dependencies: #127842
22 lines
557 B
Python
22 lines
557 B
Python
# mypy: allow-untyped-defs
|
|
# Copyright (c) Meta Platforms, Inc. and affiliates
|
|
import torch
|
|
|
|
|
|
def friendly_debug_info(v):
|
|
"""
|
|
Helper function to print out debug info in a friendly way.
|
|
"""
|
|
if isinstance(v, torch.Tensor):
|
|
return f"Tensor({v.shape}, grad={v.requires_grad}, dtype={v.dtype})"
|
|
else:
|
|
return str(v)
|
|
|
|
|
|
def map_debug_info(a):
|
|
"""
|
|
Helper function to apply `friendly_debug_info` to items in `a`.
|
|
`a` may be a list, tuple, or dict.
|
|
"""
|
|
return torch.fx.node.map_aggregate(a, friendly_debug_info)
|