问题描述

torch.cuda.memory._snapshot() 与 ._dump_snapshot() 都带有 augment_with_fx_traces
入参，torch_npu 的同名函数没有，按 CUDA 侧写法调用会直接报错。

复现步骤

import torch, torch_npu
torch.npu.memory._snapshot(augment_with_fx_traces=True)
# TypeError: _snapshot() got an unexpected keyword argument 'augment_with_fx_traces'

预期行为

与 torch.cuda.memory._snapshot() 对齐：接受该入参，为 True 时增强快照中来自 FX
生成代码的栈帧信息。

补充

社区 PR https://github.com/pytorch/pytorch/pull/187205 会让
torch.accelerator.memory._snapshot() 经 torch.get_device_module() 分发，而该调用
始终以关键字方式传递 augment_with_fx_traces。该 PR 合入后，NPU 上通过这一设备无关
接口取快照也会报同样的错。