函数功能

在图执行过程中，打印执行脚本中目标tensor值。该接口类似原生的print接口，但不会导致断图，详细功能介绍参见图内Tensor打印功能。

函数原型

def npu_print(*args, summarize_size=3)

参数说明

参数

输入/输出

说明

是否必选

*args

输入

位置入参，其中每个入参支持的数据类型为torch.Tensor、str、bool、float、int等基本类型，且其中至少包含一个Tensor类型输入。

是

summarize_size

输入

tensor每个维度显示的数据元素个数，默认值为3。

取值为-1：打印全部数据元素。
取值为正整数且大于等于（tensor维度的最大dim/2）时：打印全部数据元素。
取值为正整数且小于（tensor维度的最大dim/2）时：对于tensor每个维度，起始位置打印summarize_size个数据元素，末端位置打印summarize_size个数据元素，中间元素以“...”表示。

例如shape为[10, 10]的tensor，当summarize_size=3，打印结果示例如下：

[[1 2 3 …… 8 9 10]
[1 2 3 …… 8 9 10]
[1 2 3 …… 8 9 10]
……
[1 2 3 …… 8 9 10]
[1 2 3 …… 8 9 10]
[1 2 3 …… 8 9 10]]

否

返回值说明

无

约束说明

打印为异步打印，打印输出顺序与图中执行顺序一致，与非图内执行的其他输出顺序无关。
例如图执行后，又在外部调用Python的print打印函数，可能出现图中打印位于Python print之后打印的情况。
接口为异步接口，打印数据会占用额外的Device内存同时耗费Device侧执行时间，设置过大的summarize_size可能会因内存不足或执行超时导致失败。
请合理设置summarize_size取值，通常适合打印观察的数据量应该在KB级别以下。

例如shape为[10, 1024, 1024, 1024]的tensor，若设置summarize_size=3，内存占用量为KB级别；若设置summarize_size=1024，内存占用量会超过10G。

调用示例

import torch
import torch_npu, torchair

@torch.compile(backend="npu", fullgraph=True)
def hello_tensor(x):
    torchair.ops.npu_print("hello, tensor:", x)

v = torch.arange(10).npu()
hello_tensor(v)                          # 打印结果为"hello, tensor: [0 1 2 ... 7 8 9]"