File size: 630 Bytes
f4cade0 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 |
# Copyright (c) Meta Platforms, Inc. and affiliates
import torch
from torch.fx.node import Argument
def friendly_debug_info(v: object) -> Argument:
"""
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: Argument) -> Argument:
"""
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)
|