tvm.relax.op

Relax 核心运算符。

tvm.relax.op.assert_op(condition: RelaxExpr | PrimExpr, format_args: RelaxExpr | List[RelaxExpr] | None = None, format: str | RelaxExpr = '') → RelaxExpr

创建对 Relax 的 assert_op 操作的调用（assert 在 Python 中是保留字，因此名称必须不同）。

参数:

condition (Union[Expr, PrimExpr]) – 断言条件。
format_args (Optional[Union[Expr, List[Expr]]]) – 如果条件失败，则用于错误消息的格式化参数。
format (Union[str, Expr]) – 错误消息的格式字符串或 StringImm。

返回:

result – 对 Relax assert 操作的 relax.Call。

返回类型:

Expr

tvm.relax.op.call_builtin_with_ctx(func: str | RelaxExpr, args: RelaxExpr, *, sinfo_args: StructInfo | List[StructInfo] | None = None) → Call

relax.Call 调用内置函数 func。

参数:

func (Expr) – 要调用的内置函数。
args (Expr) – 输入参数。
sinfo_args (Optional[Union[StructInfo, List[StructInfo]]]) – 调用节点的结构信息参数。

返回:

ret – 创建的调用节点。

返回类型:

relax.Call

tvm.relax.op.call_dps_packed(func: str | RelaxExpr, args: RelaxExpr, out_sinfo: TensorStructInfo | List[TensorStructInfo]) → Call

relax.Call 调用目标传递风格的 packed 函数并返回输出。

注意：假定调用的函数是 _pure_ 的（除了修改指定的输出参数）。如果该函数 _确实_ 导致其他副作用，则编译器最终可能会删除、重新排序或重复这些影响——无法做出任何保证。

参数:

func (Union[str, Expr]) – 目标传递风格的函数，可以是 ExternFunc。
args (Expr) – 输入参数。
out_sinfo (Union[TensorStructInfo, List[TensorStructInfo]]) – call_dps_packed 输出的结构信息。它应该是单个或 TensorStructInfo 列表。每个都表示返回张量的结构信息。

返回:

ret – call_dps_packed 运算符的调用节点。

返回类型:

relax.Call

tvm.relax.op.call_inplace_packed(func: str | ExternFunc | GlobalVar, *args: RelaxExpr, inplace_indices: int | List[int], sinfo_args: StructInfo | List[StructInfo]) → RelaxExpr

构造对 packed 函数的调用，该函数“就地”消耗其某些参数并返回已变异的参数（别名），但应被认为是纯函数。inplace_indices 参数指示哪些输出是已变异的参数。

生成的调用将具有与直接调用 packed 函数相同的语义。

注意：这应该用于用户知道使用这些参数调用 packed 函数实际上不会导致任何其他副作用的情况。如果它用于确实导致其他副作用的调用，则编译器最终可能会删除、重新排序或重复该调用，并且不对被调用者的任何副作用做出任何保证。

警告：即使此运算符正在执行副作用（变异某些参数），类型系统也会将其视为纯运算符。因此，用户有责任确保安全地使用它（即，变异的参数在变异后不再有效，它们不会别名变异后仍然有效的值）。

参数:

func (Union[str, ExternFunc]) – PackedFunc 或 ExternFunc 节点的名称（全局符号）。
args (Expr) – PackedFunc 的参数。
inplace_indices (Union[int, List[int]]) – 指定应将哪些参数用于就地计算。如果 inplace_indices 是单个整数，则将其设为单例列表。假设 inplace_indices[i] = j，其中 j >= 0。那么第 i 个输出将是 args[j] 的别名。如果 inplace_indices[i] = -1，则第 i 个输出将是新分配的张量。inplace_indices 的至少一个成员不能为 -1。
sinfo_args (Union[StructInfo, List[StructInfo]]) – 结构信息参数列表（给出返回值的结构信息）。

返回:

result – Relax 调用，对应于 call_pure_packed(ExternFunc(func), args, DictAttrs(kwargs), sinfo_args)

返回类型:

Expr

tvm.relax.op.call_pure_packed(func: str | ExternFunc | GlobalVar, *args: RelaxExpr, sinfo_args: StructInfo | List[StructInfo]) → RelaxExpr

构造对 packed 函数的调用，即使 packed 调用通常不被视为纯函数，也应将其视为纯函数。

生成的调用将具有与直接调用 packed 函数相同的语义。

注意：这应该用于用户知道使用这些参数调用 packed 函数实际上不会导致任何副作用的情况。如果它用于确实导致副作用的调用，则编译器最终可能会删除、重新排序或重复该调用，并且不对被调用者的任何副作用做出任何保证。

参数:

func (Union[str, ExternFunc]) – PackedFunc 或 ExternFunc 节点的名称（全局符号）。
args (Expr) – PackedFunc 的参数。
sinfo_args (Union[StructInfo, List[StructInfo]]) – 结构信息参数列表（给出返回值的结构信息）。

返回:

result – Relax 调用，对应于 call_pure_packed(ExternFunc(func), args, DictAttrs(kwargs), sinfo_args)

返回类型:

Expr

tvm.relax.op.call_tir(gvar: GlobalVar, args: RelaxExpr, out_sinfo: TensorStructInfo | List[TensorStructInfo], tir_vars: ShapeExpr | Tuple[PrimExpr] | List[PrimExpr] | None = None) → Call

relax.Call 调用 tir.prim_func 并返回输出。

参数:

gvar (GlobalVar) – 指向 tir PrimFunc 的 GlobalVar。
args (Expr) – 输入参数。
out_sinfo (Union[TensorStructInfo, List[TensorStructInfo]]) – call_tir 输出的结构信息。它应该是单个或 TensorStructInfo 列表。每个都表示返回张量的结构信息。
tir_vars (Optional[Union[ShapeExpr, Tuple[PrimExpr], List[PrimExpr]]]) – ShapeExpr，表示调用 func 时要解包的整数元组。如果不使用，则为空

返回:

ret – call_tir 运算符的调用节点。

返回类型:

relax.Call

relax.Call 调用 TIR PrimFunc 并返回结果，根据 inplace_indices 参数就地执行指定的计算；输出将别名由就地索引选择的输入。

警告：此运算符被类型系统视为纯运算符，但实际上会变异 inplace_indices 指定的参数。不应直接使用此运算符，而应由已检查就地执行操作是否安全的 pass 插入此运算符（即，指定为输出的参数均未被别名，并且在调用 call_tir_inplace 后不再有效）。

仅应出于测试目的直接调用此运算符。

参数:

gvar (GlobalVar) – 指向 TIR PrimFunc 的 GlobalVar。
args (Expr) – 输入参数。
inplace_indices (Union[int, List[int]]) – 指定应将哪些参数用于就地计算。如果 inplace_indices 是单个整数，则将其设为单例列表。假设 inplace_indices[i] = j，其中 j >= 0。那么第 i 个输出将是 args[j] 的别名。如果 inplace_indices[i] = -1，则第 i 个输出将是新分配的张量。inplace_indices 的至少一个成员不能为 -1。
out_sinfo (Union[TensorStructInfo, List[TensorStructInfo]]) – call_tir_inplace 输出的结构信息。它应该是单个 TensorStructInfo 或 TensorStructInfo 列表。每个都表示返回张量的结构信息。如果给定 TensorStructInfo 列表，则结果将是 TensorStructInfo 元组。
tir_vars (Optional[Union[ShapeExpr, Tuple[PrimExpr], List[PrimExpr]]]) – ShapeExpr，表示调用 func 时要解包的整数元组。如果不使用，则为空

返回:

ret – call_tir 运算符的调用节点。

返回类型:

relax.Call

tvm.relax.op.call_tir_with_grad(gvar: GlobalVar, args: RelaxExpr, out_sinfo: TensorStructInfo | List[TensorStructInfo], te_grad_name: str, te_grad_kwargs: Dict[str, Object] = None, tir_vars: ShapeExpr | Tuple[PrimExpr] | List[PrimExpr] | None = None) → Call

relax.Call 一个 tir.prim_func 并返回输出。此内联函数将一个 te 梯度函数（由 te_grad_name 引用）绑定到 call_tir_with_grad 节点。te 梯度函数将由 Gradient pass 调用。

参数:

gvar (GlobalVar) – 指向 tir PrimFunc 的 GlobalVar。
args (Expr) – 输入参数。
out_sinfo (Union[TensorStructInfo, List[TensorStructInfo]]) – call_tir_with_grad 输出的结构信息。它应该是一个 TensorStructInfo 或 TensorStructInfo 列表。每一个都表示返回张量的结构信息。
te_grad_name (str) – 与 call_tir_with_grad 节点关联的 te 梯度函数的注册名称。必须作为关键字参数提供。
te_grad_kwargs (Dict[str, Object], optional) – 传递给 te 梯度函数的关键字参数。可选地作为关键字参数提供。默认值：{}。
tir_vars (Optional[Union[ShapeExpr, Tuple[PrimExpr], List[PrimExpr]]]) – ShapeExpr，表示调用 func 时要解包的整数元组。如果不使用，则为空

返回:

ret – call_tir_with_grad 运算符的调用节点。

返回类型:

relax.Call

tvm.relax.op.hint_on_device(data, dst_vdevice) → RelaxExpr

它提供了一个提示，指定输入数据应在其上执行的设备。 RealizeVDevice 利用此提示来传播虚拟设备。”

参数:

data (Expr) – 要复制的张量。
dst_device (VDevice) – 数据应该在其上执行的目标设备。

返回:

result – 结果。

返回类型:

Expr

tvm.relax.op.invoke_closure(closure: RelaxExpr, args: RelaxExpr, sinfo_args: List[StructInfo] | StructInfo) → Call

调用闭包。

参数:

closure (Expr) – VMClosure 对象。
args (Expr) – 输入参数。
type_args (Union[List[StructInfo], StructInfo]) – CallNode 的结构信息参数

返回:

ret – 对 invoke_closure 的调用。

返回类型:

relax.Call

tvm.relax.op.invoke_pure_closure(closure: RelaxExpr, args: RelaxExpr, sinfo_args: List[StructInfo] | StructInfo) → Call

调用闭包，并向编译器指示它是纯函数。

注意：这应该用于用户知道使用这些参数调用闭包**实际上**不会导致任何副作用的情况。如果它用于_确实_导致副作用的调用，那么编译器最终可能会删除、重新排序或重复该调用，并且不保证来自被调用者的任何副作用。

参数:

closure (Expr) – VMClosure 对象。
args (Expr) – 输入参数。
type_args (Union[List[StructInfo], StructInfo]) – CallNode 的结构信息参数

返回:

ret – 对 invoke_pure_closure 的调用。

返回类型:

relax.Call

tvm.relax.op.make_closure(func: RelaxExpr, args: RelaxExpr) → Object

创建带有自由变量的闭包并返回该闭包。

参数:

func (Expr) – 闭包，可以是 ExternFunc 或 PrimFunc。
args (Expr) – 输入参数。

返回:

ret – VMClosure。

返回类型:

Object

tvm.relax.op.null_value() → Call

创建一个表示空值对象的调用节点。

返回:: ret – 创建的调用节点。
返回类型:: relax.Call

tvm.relax.op.print(*values: List[RelaxExpr], format: str | RelaxExpr = '') → RelaxExpr

打印操作以打印值

参数:

values (List[Expr]) – 要打印的值。
format (Union[str, Expr]) – 格式字符串或 StringImm。

返回:

result – 一个 relax relax.Call，它将在运行时打印值。

返回类型:

Expr

tvm.relax.op.register_gradient(op_name: str, fgradient: Callable[[Var, Call, Var, BlockBuilder], List[RelaxExpr]] = None, level: int = 10)

为 relax 运算符注册运算符梯度函数。

参数:

op_name (str) – 运算符的名称。
fgradient (function (orig_var: relax.Var, orig_call: relax.Call, output_grad: relax.Var, ctx: BlockBuilder)) – -> partials: List[Expr] 正在使用的梯度函数。
level (int) – 优先级

tvm.relax.op.shape_of(expr: RelaxExpr) → RelaxExpr

获取张量的形状。

参数:: expr (Expr) – 输入 Expr。
返回:: result – 一个 relax relax.Call，它获取输入的形状
返回类型:: Expr

tvm.relax.op.shape_to_tensor(expr: RelaxExpr) → RelaxExpr

将形状转换为张量表达式。 :param expr: 输入 Expr :type expr: Expr

返回:: result – 一个 relax relax.Call，它将形状值转换为张量
返回类型:: Expr

tvm.relax.op.tensor_to_shape(expr: RelaxExpr) → RelaxExpr

将张量转换为形状表达式。 :param expr: 输入 Expr :type expr: Expr

返回:: result – 一个 relax relax.Call，它将张量值转换为形状
返回类型:: Expr

tvm.relax.op.to_vdevice(data, dst_vdevice) → RelaxExpr

将数据复制到目标设备。此运算符有助于为异构执行在不同设备之间传输数据。

参数:

data (Expr) – 要复制的张量。
dst_device (VDevice) – 数据复制到的目标设备。

返回:

result – 复制的结果。

返回类型:

Expr

tvm.relax.op.add(x1: RelaxExpr, x2: RelaxExpr) → RelaxExpr

使用 numpy 风格广播进行加法。

参数:

x1 (Expr) – 第一个输入张量。
x2 (Expr) – 第二个输入张量。

返回:

result – 计算结果。

返回类型:

Expr

示例

bb = relax.BlockBuilder()
a = relax.Var("a", relax.TensorStructInfo(shape=(2, 3), dtype="float32"))
b = relax.Var("b", relax.TensorStructInfo(shape=(2, 1), dtype="float32"))
c = bb.normalize(relax.op.add(a, b))  # c has TensorStructInfo(shape=(2, 3), dtype="float32")

tvm.relax.op.bitwise_and(x1: RelaxExpr, x2: RelaxExpr) → RelaxExpr

按位与 :param x1: 第一个输入张量。 :type x1: relax.Expr :param x2: 第二个输入张量。 :type x2: relax.Expr

返回:: result – 计算结果。
返回类型:: relax.Expr

tvm.relax.op.bitwise_or(x1: RelaxExpr, x2: RelaxExpr) → RelaxExpr

按位或 :param x1: 第一个输入张量。 :type x1: relax.Expr :param x2: 第二个输入张量。 :type x2: relax.Expr

返回:: result – 计算结果。
返回类型:: relax.Expr

tvm.relax.op.bitwise_xor(x1: RelaxExpr, x2: RelaxExpr) → RelaxExpr

按位异或 :param x1: 第一个输入张量。 :type x1: relax.Expr :param x2: 第二个输入张量。 :type x2: relax.Expr

返回:: result – 计算结果。
返回类型:: relax.Expr

tvm.relax.op.divide(x1: RelaxExpr, x2: RelaxExpr) → RelaxExpr

使用 numpy 风格广播进行除法。

参数:

x1 (relax.Expr) – 第一个输入张量。
x2 (relax.Expr) – 第二个输入张量。

返回:

result – 计算结果。

返回类型:

relax.Expr

tvm.relax.op.equal(x1: RelaxExpr, x2: RelaxExpr) → RelaxExpr

广播元素方式测试 (lhs == rhs)。

参数:

x1 (relax.Expr) – 第一个输入张量。
x2 (relax.Expr) – 第二个输入张量。

返回:

result – 计算结果。

返回类型:

relax.Expr

tvm.relax.op.floor_divide(x1: RelaxExpr, x2: RelaxExpr) → RelaxExpr

使用 numpy 风格广播进行向下取整除法。

参数:

x1 (relax.Expr) – 第一个输入张量。
x2 (relax.Expr) – 第二个输入张量。

返回:

result – 计算结果。

返回类型:

relax.Expr

tvm.relax.op.floor_mod(x1: RelaxExpr, x2: RelaxExpr) → RelaxExpr

使用 numpy 风格广播进行向下取整求模。

参数:

x1 (Expr) – 第一个输入张量。
x2 (Expr) – 第二个输入张量。

tvm.relax.op.greater(x1: RelaxExpr, x2: RelaxExpr) → RelaxExpr

广播元素方式测试 (lhs > rhs)。

参数:

x1 (relax.Expr) – 第一个输入张量。
x2 (relax.Expr) – 第二个输入张量。

返回:

result – 计算结果。

返回类型:

relax.Expr

tvm.relax.op.greater_equal(x1: RelaxExpr, x2: RelaxExpr) → RelaxExpr

广播元素方式测试 (lhs >= rhs)。

参数:

x1 (relax.Expr) – 第一个输入张量。
x2 (relax.Expr) – 第二个输入张量。

返回:

result – 计算结果。

返回类型:

relax.Expr

tvm.relax.op.left_shift(x1: RelaxExpr, x2: RelaxExpr) → RelaxExpr

按位左移 :param x1: 要移位的输入张量。 :type x1: relax.Expr :param x2: 要移位的位数。 :type x2: relax.Expr

返回:: result – 计算结果。
返回类型:: relax.Expr

tvm.relax.op.less(x1: RelaxExpr, x2: RelaxExpr) → RelaxExpr

广播元素方式测试 (lhs < rhs)。

参数:

x1 (relax.Expr) – 第一个输入张量。
x2 (relax.Expr) – 第二个输入张量。

返回:

result – 计算结果。

返回类型:

relax.Expr

tvm.relax.op.less_equal(x1: RelaxExpr, x2: RelaxExpr) → RelaxExpr

广播元素方式测试 (lhs <= rhs)。

参数:

x1 (relax.Expr) – 第一个输入张量。
x2 (relax.Expr) – 第二个输入张量。

返回:

result – 计算结果。

返回类型:

relax.Expr

tvm.relax.op.logical_and(x1: RelaxExpr, x2: RelaxExpr) → RelaxExpr

逻辑与 :param x1: 第一个输入张量。 :type x1: relax.Expr :param x2: 第二个输入张量。 :type x2: relax.Expr

返回:: result – 计算结果。
返回类型:: relax.Expr

tvm.relax.op.logical_or(x1: RelaxExpr, x2: RelaxExpr) → RelaxExpr

逻辑或 :param x1: 第一个输入张量。 :type x1: relax.Expr :param x2: 第二个输入张量。 :type x2: relax.Expr

返回:: result – 计算结果。
返回类型:: relax.Expr

tvm.relax.op.logical_xor(x1: RelaxExpr, x2: RelaxExpr) → RelaxExpr

逻辑异或 :param x1: 第一个输入张量。 :type x1: relax.Expr :param x2: 第二个输入张量。 :type x2: relax.Expr

返回:: result – 计算结果。
返回类型:: relax.Expr

tvm.relax.op.maximum(x1: RelaxExpr, x2: RelaxExpr) → RelaxExpr

元素方式最大值

参数:

x1 (relax.Expr) – 第一个输入张量。
x2 (relax.Expr) – 第二个输入张量。

返回:

result – 计算结果。

返回类型:

relax.Expr

tvm.relax.op.minimum(x1: RelaxExpr, x2: RelaxExpr) → RelaxExpr

元素方式最小值

参数:

x1 (relax.Expr) – 第一个输入张量。
x2 (relax.Expr) – 第二个输入张量。

返回:

result – 计算结果。

返回类型:

relax.Expr

tvm.relax.op.mod(x1: RelaxExpr, x2: RelaxExpr) → RelaxExpr

NumPy 风格的广播取模。

参数:

x1 (Expr) – 第一个输入张量。
x2 (Expr) – 第二个输入张量。

tvm.relax.op.multiply(x1: RelaxExpr, x2: RelaxExpr) → RelaxExpr

NumPy 风格的广播乘法。

参数:

x1 (Expr) – 第一个输入张量。
x2 (Expr) – 第二个输入张量。

返回:

result – 计算结果。

返回类型:

Expr

tvm.relax.op.not_equal(x1: RelaxExpr, x2: RelaxExpr) → RelaxExpr

广播元素级不等式测试 (lhs != rhs)。

参数:

x1 (relax.Expr) – 第一个输入张量。
x2 (relax.Expr) – 第二个输入张量。

返回:

result – 计算结果。

返回类型:

relax.Expr

tvm.relax.op.power(x1: RelaxExpr, x2: RelaxExpr)

NumPy 风格的广播幂运算。

参数:

x1 (relax.Expr) – 第一个输入张量。
x2 (relax.Expr) – 第二个输入张量。

返回:

result – 计算结果。

返回类型:

relax.Expr

tvm.relax.op.right_shift(x1: RelaxExpr, x2: RelaxExpr) → RelaxExpr

位右移 :param x1: 要进行位移的输入张量。 :type x1: relax.Expr :param x2: 位移的位数。 :type x2: relax.Expr

返回:: result – 计算结果。
返回类型:: relax.Expr

tvm.relax.op.subtract(x1: RelaxExpr, x2: RelaxExpr) → RelaxExpr

NumPy 风格的广播减法。

参数:

x1 (relax.Expr) – 第一个输入张量。
x2 (relax.Expr) – 第二个输入张量。

返回:

result – 计算结果。

返回类型:

relax.Expr

构造一个包含均匀间隔元素的张量。

参数:

start (Union[PrimExprLike,PrimValue]) – 区间的起始值。
end (Optional[Union[PrimExprLike,PrimValue]]) – 区间的结束值。如果未给定，则设置为 start，并且 start 将设置为 0。
step (Union[PrimExprLike,PrimValue]) – 步长。
dtype (Optional[Union[str, DataType]]) – 创建的张量的数据类型。

返回:

result – 结果张量。

返回类型:

relax.Expr

tvm.relax.op.full(shape: Tuple[int | PrimExpr] | RelaxExpr, fill_value: RelaxExpr, dtype: str | DataType | None = None) → RelaxExpr

使用标量值填充数组。

参数:

shape (Union[Tuple[PrimExprLike], Expr]) – 创建的张量的形状。
fill_value (relax.Expr) – 要填充的值。必须是标量张量。
dtype (Optional[Union[str, DataType]]) – 创建的张量的数据类型。如果未给定 dtype，默认情况下将使用 fill_value 的 dtype。

返回:

result – 结果张量。

返回类型:

relax.Expr

tvm.relax.op.full_like(x: RelaxExpr, fill_value: RelaxExpr, dtype: str | DataType | None = None) → RelaxExpr

构造一个张量，使其： - 形状与输入数据张量的形状相同， - 其值填充为输入的标量填充值。

参数:

x (relax.Expr) – 输入张量，提供形状，并在未指定 dtype 字段时提供 dtype。
fill_value (relax.Expr) – 要填充的值。必须是标量张量。
dtype (Optional[Union[str, DataType]]) – 创建的张量的数据类型。如果未给定 dtype，默认情况下将使用输入张量的 dtype。

返回:

result – 结果张量。

返回类型:

relax.Expr

tvm.relax.op.ones(shape: Tuple[int | PrimExpr] | RelaxExpr, dtype: str | DataType) → RelaxExpr

构造一个所有元素都为 1 的张量，具有输入形状和 dtype。

参数:

shape (Union[Tuple[PrimExprLike], Expr]) – 创建的张量的形状。
dtype (Union[str, DataType]) – 创建的张量的数据类型。

返回:

result – 结果张量。

返回类型:

relax.Expr

tvm.relax.op.ones_like(x: RelaxExpr, dtype: str | DataType | None = None) → RelaxExpr

构造一个所有元素都为 1 的张量，其形状与输入张量的形状相同。

参数:

x (relax.Expr) – 输入张量，提供形状，并在未指定 dtype 字段时提供 dtype。
dtype (Optional[Union[str, DataType]]) – 创建的张量的数据类型。如果未给定 dtype，默认情况下将使用输入张量的 dtype。

返回:

result – 结果张量。

返回类型:

relax.Expr

构造一个 2 维张量，对角线上为 1，其他地方为 0。

参数:

n (Union[PrimExprLike, PrimValue]) – 输出中的行数。
m (Optional[Union[PrimExprLike, PrimValue]]) – 输出中的列数。如果为 None，则默认为 n。
k (Union[PrimExprLike, PrimValue]) – 对角线的索引：0（默认值）指主对角线，正值指上对角线，负值指下对角线。
dtype (Union[str, DataType]) – 创建的张量的数据类型。

返回:

result – 结果张量。

返回类型:

relax.Expr

tvm.relax.op.eye_like(x: RelaxExpr, k: int | PrimExpr | PrimValue = 0, dtype: str | DataType | None = None) → RelaxExpr

返回一个 2 维张量，对角线上为 1，其他地方为 0，形状与输入张量相同。

参数:

x (relax.Expr) – 输入张量，提供形状，并在未指定 dtype 字段时提供 dtype。
k (Union[PrimExprLike, PrimValue]) – 对角线的索引：0（默认值）指主对角线，正值指上对角线，负值指下对角线。
dtype (Optional[Union[str, DataType]]) – 创建的张量的数据类型。如果未给定 dtype，默认情况下将使用输入张量的 dtype。

返回:

result – 结果张量。

返回类型:

relax.Expr

tvm.relax.op.tril(x: RelaxExpr, k: int | PrimExpr | RelaxExpr = 0) → RelaxExpr

返回矩阵或一批矩阵的下三角部分。

参数:

x (relax.Expr) – 将应用 tril 的张量。要求至少具有两个维度。
k (int) – 指示对角线索引，高于该对角线的元素将置零。如果 k = 0，则对角线为主对角线。如果 k < 0，则对角线位于主对角线下方。如果 k > 0，则对角线位于主对角线上方。

返回:

ret – 结果张量。

返回类型:

relax.Expr

tvm.relax.op.triu(x: RelaxExpr, k: int | PrimExpr | RelaxExpr = 0) → RelaxExpr

返回矩阵或一批矩阵的上三角部分。

参数:

x (relax.Expr) – 将应用 triu 的张量。要求至少具有两个维度。
k (int) – 指示对角线索引，低于该对角线的元素将置零。如果 k = 0，则对角线为主对角线。如果 k < 0，则对角线位于主对角线下方。如果 k > 0，则对角线位于主对角线上方。

返回:

ret – 结果张量。

返回类型:

relax.Expr

tvm.relax.op.zeros(shape: Tuple[int | PrimExpr] | RelaxExpr, dtype: str | DataType) → RelaxExpr

构造一个所有元素都为 0 的张量，具有输入形状和 dtype。

参数:

shape (Union[Tuple[PrimExprLike], Expr]) – 创建的张量的形状。
dtype (Union[str, DataType]) – 创建的张量的数据类型。

返回:

result – 结果张量。

返回类型:

relax.Expr

tvm.relax.op.zeros_like(x: RelaxExpr, dtype: str | DataType | None = None) → RelaxExpr

构造一个所有元素都为 0 的张量，其形状与输入张量的形状相同。

参数:

x (relax.Expr) – 输入张量，提供形状，并在未指定 dtype 字段时提供 dtype。
dtype (Optional[Union[str, DataType]]) – 创建的张量的数据类型。如果未给定 dtype，默认情况下将使用输入张量的 dtype。

返回:

result – 结果张量。

返回类型:

relax.Expr

tvm.relax.op.astype(x: RelaxExpr, dtype: str | DataType) → RelaxExpr

将输入张量转换为给定的数据类型。

参数:

x (relax.Expr) – 运算符的输入数据。
dtype (Union[str, DataType]) – 目标数据类型

返回:

result – 转换后的结果。

返回类型:

relax.Expr

tvm.relax.op.wrap_param(data: RelaxExpr, dtype: str | DataType = 'float32') → RelaxExpr

如果输入数据的dtype与给定的dtype不同，则将作为模型参数的输入张量转换为数据类型。 :param data: 运算符的输入数据。 :type data: relax.Expr :param dtype: 目标数据类型 :type dtype: Union[str, DataType]

返回:: result – 转换后的结果。
返回类型:: relax.Expr

tvm.relax.op.dynamic_strided_slice(x: RelaxExpr, begin: RelaxExpr, end: RelaxExpr, strides: RelaxExpr) → RelaxExpr

张量的动态步幅切片。 begin，end，strides 可以在运行时计算。

参数:

x (Expr) – 要切片的源张量。
begin (Expr) – 切片开始的索引，包含在内。
end (Expr) – 指示切片结束的索引，不包含在内。
strides (Expr) – 指定步幅值，它可以为负数，在这种情况下，输入张量将在该特定轴中反转。如果未指定，则默认情况下它是与 axes 长度相同的 1 列表。

返回:

ret – 切片结果。

返回类型:

relax.Expr

Note

dyn_strided_slice 要求输入 begin，end 和 strides 的长度与 data 张量的秩相同。

tvm.relax.op.strided_slice(x: RelaxExpr, axes: RelaxExpr, begin: RelaxExpr, end: RelaxExpr, strides: RelaxExpr | None = None, assume_inbound: bool = False) → RelaxExpr

张量的步幅切片。

参数:

x (relax.Expr) – 要切片的源张量。
axes (List[int]) – 应用切片的轴。
begin (List[PrimExprLike]) – 切片开始的索引，包含在内。
end (List[PrimExprLike]) – 指示切片结束的索引，不包含在内。
strides (Optional[List[PrimExprLike]]) – 指定步幅值，它可以为负数，在这种情况下，输入张量将在该特定轴中反转。如果未指定，则默认情况下它是与 axes 长度相同的 1 列表。
assume_inbound (bool) – 是否假定索引在界内。如果设置为 false，则超出界限的索引将被裁剪到界限内。

返回:

ret – 切片结果。

返回类型:

relax.Expr

Note

strided_slice 要求输入 begin，end 和 strides 的长度与 axes 相同。

tvm.relax.op.take(x: RelaxExpr, indices: RelaxExpr, axis: int | None = None) → RelaxExpr

沿轴从张量中提取元素。它的语义主要类似于 numpy.take (https://numpy.com.cn/doc/stable/reference/generated/numpy.take.html)，它可以覆盖 torch.take (https://pytorch.ac.cn/docs/stable/generated/torch.take.html) 和 onnx.gather (https://github.com/onnx/onnx/blob/main/docs/Changelog.md#Gather-13)。

参数:

x (relax.Expr) – 源张量。
indices (relax.Expr) – 要提取的值的索引。
axis (Optional[int]) – 选择值的轴。如果为 None，则输入张量必须是一维的。

返回:

ret – 提取的结果。

返回类型:

relax.Expr

tvm.relax.op.einsum(operands, subscripts)

在数据上评估爱因斯坦求和约定

参数:

operands (Union(List[relax.Expr], Tuple[relax.Expr])) – 表达式列表。
subscripts (str) – einsum 表达式字符串。

返回:

result – 来自 einsum 运算的输出。

返回类型:

relax.Expr

tvm.relax.op.linear(data: RelaxExpr, weight: RelaxExpr, bias: RelaxExpr | None = None, out_dtype: str | DataType | None = None) → RelaxExpr

对传入数据应用线性变换：y = xA^T + b

参数:

data (relax.Expr) – 输入数据。
weight (relax.Expr) – 权重张量。
bias (Optional[Expr]) – 偏置张量。
out_dtype (Optional[Union[str, DataType]]) – matmul 结果的数据类型。如果未指定，则输出 dtype 将与输入 dtype 相同。

Notes

Relax 不将 Linear Op 视为原始 Op，而是结合转置、matmul 和 add op 来实现它。

返回:: result – 计算结果。
返回类型:: relax.Expr

tvm.relax.op.matmul(x1: RelaxExpr, x2: RelaxExpr, out_dtype: str | DataType | None = None) → RelaxExpr

两个张量的一般矩阵乘法，在批次维度上进行广播。

语义和输出形状推导规则指定为 https://data-apis.org/array-api/latest/API_specification/generated/array_api.matmul.html。

参数:

x1 (relax.Expr) – 第一个输入张量。
x2 (relax.Expr) – 第二个输入张量。
out_dtype (Optional[Union[str, DataType]]) – matmul 结果的数据类型。如果未指定，则输出 dtype 将与输入 dtype 相同。

返回:

result – 计算结果。

返回类型:

relax.Expr

tvm.relax.op.broadcast_to(x: RelaxExpr, shape: Tuple[int | PrimExpr] | RelaxExpr) → RelaxExpr

将张量广播到指定的形状。

参数:

x (relax.Expr) – 运算符的输入数据。
shape (Union[Tuple[PrimExprLike], Expr]) – 目标形状。

返回:

result – 广播后的张量。

返回类型:

relax.Expr

tvm.relax.op.collapse_sum_like(data: RelaxExpr, collapse_target: RelaxExpr) → RelaxExpr

返回与 collapse_target 形状相同的 data 的总和。

有关详细信息，请参阅 relax.op.collapse_sum_to。

参数:

data (relax.Expr) – 输入张量。
collapse_target (relax.Expr) – 其形状是要折叠到的形状的张量。

返回:

result – 求和后的结果张量。

返回类型:

relax.Expr

tvm.relax.op.collapse_sum_to(data: RelaxExpr, shape: Tuple[int | PrimExpr] | RelaxExpr) → RelaxExpr

返回与给定形状相同的 data 的总和。

collapse_sum_to 旨在作为 tvm.relax.op.broadcast_to 和自动微分过程中的其他广播运算符的反向运算符。

我们期望 data 是在某些广播操作中广播给定形状的某个张量的结果。因此，给定的 shape 和 data.shape 必须遵循广播规则。

在计算期间，从右到左检查 data.shape 和 shape 的所有轴。对于一个轴，如果它遵循以下规则，则 data 将在该轴上求和： - 该轴存在于 data.shape 中，但不存在于 shape 中，或者 - 该轴存在于 data.shape 中，并且在 shape 中等于 1。

参数:

data (relax.Expr) – 输入张量。
shape (Union[Tuple[PrimExprLike], relax.Expr]) – 要折叠到的形状。

返回:

result – 求和后给定形状的结果张量。

返回类型:

relax.Expr

tvm.relax.op.concat(tensors: RelaxExpr | List[RelaxExpr], axis: int | None = 0) → RelaxExpr

沿给定轴连接输入张量。

参数:

tensors (Union[relax.Expr, List[relax.Expr]]) – Tuple 类型中的 Expr，包含要连接的张量，或张量列表。
axis (Optional[int]) – 张量连接的轴。如果 axis 为 None，则需要先展平输入张量，然后再进行连接。

返回:

result – 连接后的张量。

返回类型:

relax.Expr

tvm.relax.op.expand_dims(x: RelaxExpr, axis: int | List[int]) → RelaxExpr

在 axis 给定的位置插入新轴。

参数:

x (relax.Expr) – 运算符的输入数据。
axis (Union[int, List[int]]) – 输入数组扩展的轴。所有值都必须在范围 [-data.ndim - 1, data.ndim] 内，并采用负索引的约定。

返回:

result – 转换后的结果。

返回类型:

relax.Expr

tvm.relax.op.flatten(x: RelaxExpr) → RelaxExpr

将所有张量维度展平为一个维度。

参数:: x (relax.Expr) – 运算符的输入数据。
返回:: result – 展平的结果。
返回类型:: relax.Expr

tvm.relax.op.flip(data, axis)

反转给定轴上元素的顺序，同时保留数组形状。

参数:

data (relax.Expr) – 运算符的输入数据。
axis (int) – 要翻转的轴

返回:

ret – 计算结果。

返回类型:

relax.Expr

示例

x = [[1., 2.], [3., 4.]]
relax.flip(x, axis=0) = [[3., 4.], [1., 2.]]

relax.flip(x, axis=1) = [[2., 1.], [4., 3.]]

tvm.relax.op.gather_elements(data: RelaxExpr, indices: RelaxExpr, axis: int = 0) → RelaxExpr

根据指定轴的索引从数据中收集元素。

参数:

data (relax.Expr) – 运算符的输入数据。
indices (relax.Expr) – 索引张量，必须为整数类型。
axis (int) – 索引的轴。默认为 0。

返回:

ret – 计算结果。

返回类型:

relax.Expr

示例

data = [[1, 2], [3, 4]]
indices = [[0, 0], [1, 0]]
axis = 1
output = [[1, 1], [4, 3]]

data = [[1, 2, 3], [4, 5, 6]]
indices = [[1, 1, 1]]
axis = 0
output = [[4, 5, 6]]

tvm.relax.op.gather_nd(data: RelaxExpr, indices: RelaxExpr, batch_dims: int = 0) → RelaxExpr

使用更新中的值更新索引定义位置的数据。

参数:

data (relax.Expr) – 运算符的输入数据。
indices (relax.Expr) – 索引张量，必须为整数类型。
batch_dims (int) – 批次维度的数量。默认为 0。

返回:

ret – 计算结果。

返回类型:

relax.Expr

示例

batch_dims = 0
data    = [[0,1],[2,3]]   # data_shape    = [2, 2]
indices = [[0,0],[1,1]]   # indices_shape = [2, 2]
output  = [0,3]           # output_shape  = [2]

batch_dims = 1
data    = [[[0,1],[2,3]],[[4,5],[6,7]]] # data_shape    = [2, 2, 2]
indices = [[1],[0]]                     # indices_shape = [2, 1]
output  = [[2,3],[4,5]]                 # output_shape  = [2, 2]

修改张量的布局。

参数:

x (relax.Expr) – 算子的输入张量。
index_map (Union[Callable, IndexMap]) – 要应用的变换。
pad_value (Optional[Union[int, float, PrimValue]]) – 如果变换导致隐式填充，则用于填充的值。如果未指定，则可以使用任何值。
axis_separators (Optional[Union[int, IndexMap.AXIS_SEPARATOR]]) – 用于 index_map 的 axis_separators，以创建非扁平缓冲区。

返回:

result – 变换后的张量。

返回类型:

relax.Expr

tvm.relax.op.one_hot(indices: RelaxExpr, on_value: PrimValue, off_value: PrimValue, depth: int, axis: int = -1) → RelaxExpr

返回一个 one-hot 张量。

参数:

indices (relax.Expr) – 要设置为 on_value 的索引。
on_value (relax.PrimValue) – 在 indices 处填充的值。
off_value (relax.PrimValue) – 在其他位置填充的值。
depth (int) – one-hot 维度的深度。
axis (int, optional) – 要填充的轴。默认为 -1，这将在末尾添加一个新维度。

返回:

result – 计算结果。

返回类型:

relax.Expr

示例

indices = [0, 1, 2]
depth = 3
on_value = 1
off_value = 0

one_hot(indices, on_value, off_value, depth) =
    [[1, 0, 0],
     [0, 1, 0],
     [0, 0, 1]]

tvm.relax.op.permute_dims(x: RelaxExpr, axes: List[int] | None = None) → RelaxExpr

置换数组的维度。

参数:

x (relax.Expr) – 运算符的输入数据。
axes (Optional[List[int]]) – 目标轴顺序。如果未指定，permute_dims 将反转所有轴的顺序。

返回:

result – 转置后的结果。

返回类型:

relax.Expr

tvm.relax.op.repeat(data: RelaxExpr, repeats: int, axis: int | None = None) → RelaxExpr

重复数组的元素。

参数:

data (relax.Expr) – 输入张量。
repeats (int) – 重复次数。
axis (Optional[int]) – 沿其重复值的轴。负数从后往前计数。默认情况下，使用扁平化的输入数组，并返回扁平化的输出数组。

返回:

ret – 计算结果。

返回类型:

relax.Expr

示例

x = R.const([[1, 2], [3, 4]])
lv1 = R.repeat(x, repeats=2) # lv1 == [1, 1, 2, 2, 3, 3, 4, 4]
lv2 = R.repeat(x, repeats=2, axis=1) # lv2 == [[1., 1., 2., 2.],
                                     #         [3., 3., 4., 4.]]

tvm.relax.op.reshape(x: RelaxExpr, shape: Tuple[int | PrimExpr] | RelaxExpr) → RelaxExpr

重塑输入数组。

-1 通过使用输入维度的剩余部分来推断输出形状的维度，保持新数组的大小与输入数组的大小相同。形状中最多可以有一个维度为 -1。

x.shape = (2, 3, 4), shape = (6, 1, -1), result.shape = (6, 1, 4)
x.shape = (2, 3, 4), shape = (3, -1, 8), result.shape = (3, 1, 8)
x.shape = (2, 3, 4), shape = (-1,), result.shape = (24,)

参数:

x (relax.Expr) – 运算符的输入数据。
shape (Union[Tuple[PrimExprLike], Expr]) – 新形状。应与原始形状兼容。

返回:

result – 重塑后的结果。

返回类型:

relax.Expr

Note

-1 推断仅在编译时执行。也就是说，在任何情况下，如果无法在编译时推断 -1 的维度长度，则会抛出错误。

tvm.relax.op.scatter_elements(data: RelaxExpr, indices: RelaxExpr, updates: RelaxExpr, axis: int = 0, reduction: str = 'update')

ONNX 风格的 scatter elements。此操作在 indices 指定的特定索引位置，将其在 data 中的值更新为 updates 指定的值。例如，在 2D 张量中，对应于 [i][j] 条目的更新执行如下：

output[indices[i][j]][j] = updates[i][j] if axis = 0
output[i][indices[i][j]] = updates[i][j] if axis = 1

当 reduction 设置为某个归约函数 f 时，对应于 [i][j] 条目的更新执行如下：

output[indices[i][j]][j] += f(output[indices[i][j]][j], updates[i][j]) if axis = 0
output[i][indices[i][j]] += f(output[i][indices[i][j]], updates[i][j]) if axis = 1

其中 f 是 update、add、mul、mean、max、min。

参数:

data (relax.Expr) – 运算符的输入数据。
indices (relax.Expr) – 要在 data 中更新的索引位置。
updates (relax.Expr) – 要替换为的值。
axis (int) – 要在其上进行 scatter 的轴。
reduction (str) – 要应用的归约类型：update、add、mul、mean、max、min。默认为 “update”。

返回:

result – 结果的大小与 data 相同，形状与 data 相同

返回类型:

relax.Expr

示例

# inputs
data = [
     [0.0, 0.0, 0.0],
     [0.0, 0.0, 0.0],
     [0.0, 0.0, 0.0],
 ]
 indices = [
     [1, 0, 2],
     [0, 2, 1],
 ]
 updates = [
     [1.0, 1.1, 1.2],
     [2.0, 2.1, 2.2],
 ]
 axis = 0
 reduction = "update"

 # output P
 output = [
     [2.0, 1.1, 0.0]
     [1.0, 0.0, 2.2]
     [0.0, 2.1, 1.2]
 ]

tvm.relax.op.scatter_nd(data: RelaxExpr, indices: RelaxExpr, updates: RelaxExpr, reduction: str = 'update') → RelaxExpr

根据索引将更新分散到数组中。

参数:

data (relax.Expr) – 要更新的输入数据。
indices (relax.Expr) – 要在 data 中更新的索引位置。
updates (relax.Expr) – 要替换为的值。
reduction (str) – 要应用的归约类型：update、add、mul、max、min。默认为 “update”。

返回:

result – 结果具有与 data 相同的形状。

返回类型:

relax.Expr

示例

# inputs
data = [1, 2, 3, 4, 5, 6, 7, 8]
indices = [[4], [3], [1], [7]]
updates = [9, 10, 11, 12]

# output
output = [1, 11, 3, 10, 9, 6, 7, 12]

tvm.relax.op.split(x: RelaxExpr, indices_or_sections: int | List[int | PrimExpr], axis: int = 0) → RelaxExpr

沿轴按节或索引拆分输入张量。

如果 indices_or_sections 是一个整数，则输入将沿给定轴均匀分割（如果可能）。如果张量沿给定维度的大小不能被整数整除，则最后一节会更小。

如果 indices_or_sections 是 int 或 PrimExpr 混合的元组，则条目指示沿轴分割数组的索引。

参数:

x (relax.Expr) – 要拆分的张量。
indices_or_sections (Union[int, List[PrimExprLike]]) – 要拆分成的索引或节。接受 int 或列表。
axis (int) – 要在其上进行拆分的轴。

返回:

ret – 计算结果。

返回类型:

relax.Expr

tvm.relax.op.squeeze(x: RelaxExpr, axis: int | List[int] | None = None) → RelaxExpr

挤压数组中的轴。

参数:

x (relax.Expr) – 运算符的输入数据。
axis (Optional[Union[int, List[int]]) – 要移除的轴的集合。如果 axis = None，则移除维度为 1 的所有轴。如果任何指定的轴的维度不等于 1，则会报错。

返回:

result – 挤压后的结果。

返回类型:

relax.Expr

tvm.relax.op.tile(data: RelaxExpr, repeats: int | Tuple[int] | List[int]) → RelaxExpr

通过将数据重复给定的次数来构造数组。

如果 repeats 的长度为 l，而数据的维度为 d，则结果的维度将为 max(l, d)。

如果 d < l，则通过预先添加新轴将数据提升为 l 维。因此，形状为 (3,) 的张量对于 2-D 复制会提升为 (1, 3)，对于 3-D 复制会提升为形状 (1, 1, 3)。如果这不是所需的行为，请在调用此函数之前手动将数据提升为 d 维。

如果 d > l，则通过在其前面添加 1 将 reps 提升为长度 d。因此，对于形状为 (2, 3, 4, 5) 的数据，(2, 2) 的 reps 被视为 (1, 1, 2, 2)。

参数:

data (relax.Expr) – 运算符的输入数据。
repeats (Union[int, Tuple[int], List[int]]) – 数据沿每个轴的重复次数。

返回:

ret – 计算结果。

返回类型:

relax.Expr

示例

x = R.const([[1, 2], [3, 4]])
lv1 = R.tile(x, reps=(2, 3)) # lv1 = [[1., 2., 1., 2., 1., 2.],
                             #        [3., 4., 3., 4., 3., 4.],
                             #        [1., 2., 1., 2., 1., 2.],
                             #        [3., 4., 3., 4., 3., 4.]]
lv2 = R.tile(x, reps=2) # lv2 = [[1., 2., 1., 2.],
                        #        [3., 4., 3., 4.]]

tvm.relax.op.masked_fill(x: RelaxExpr, mask: RelaxExpr, value: RelaxExpr)

在由掩码定义的位置用指定值填充张量。 :param x: 算子的输入数据。 :type x: relax.Expr :param mask: 掩码。 :type mask: relax.Expr :param value: 在输入张量中设置的值。 :type value: relax.Expr

返回:: result – 填充后的张量。
返回类型:: relax.Expr

tvm.relax.op.dequantize(data: RelaxExpr, scale: RelaxExpr, zero_point: RelaxExpr, axis: int = -1, out_dtype: str = 'float32')

反量化算子。此算子接受输入并生成反量化输出。输入张量可以是任何形状。输出形状与输入形状相同。

output = clamp(scale * (input_tensor - zero_point), out_dtype::min, out_dtype::max)

参数:

data (tvm.relax.Expr) – 要反量化的输入张量。
scale (tvm.relax.Expr) – 输入缩放。
zero_point (tvm.relax.Expr) – 输入零点。
axis (int) – 反量化的通道轴。默认值为 -1，对应于最后一个轴。
out_dtype (str, optional) – 输出张量的数据类型。

返回:

result – 计算结果。

返回类型:

tvm.relax.Expr

tvm.relax.op.quantize(data: RelaxExpr, scale: RelaxExpr, zero_point: RelaxExpr, axis: int = -1, out_dtype: str = 'int8')

量化算子。此算子接受输入并生成量化输出。输入张量可以是任何形状。输出形状与输入形状相同。

Q_output = clamp((round(input_tensor/scale) + zero_point), out_dtype::min, out_dtype::max)

参数:

data (tvm.relax.Expr) – 要量化的输入张量。
scale (tvm.relax.Expr) – 输出缩放。
zero_point (tvm.relax.Expr) – 输出零点。
axis (int) – 量化的通道轴。默认值为 -1，对应于最后一个轴。
out_dtype (str, optional) – 输出张量的数据类型。

返回:

result – 计算结果。

返回类型:

tvm.relax.Expr

tvm.relax.op.multinomial_from_uniform(prob: RelaxExpr, uniform_sample: RelaxExpr, sample_indices: RelaxExpr, dtype: str = 'int64') → RelaxExpr

返回一个张量，其中每行包含从张量 prob 的相应行中的多项概率分布中采样的索引。

Notes

为了获得更好的 CPU 性能，请使用 ‘vm.builtin.multinomial_from_uniform’。为了获得准确的结果，请确保概率介于 0 和 1 之间，并且总和为 1。

参数:

prob (relax.Expr) – 形状为 (batch, vocab_size) 的 2-D 张量，表示概率分布。每行是批次词汇表的分布，其中：值范围为 [0, 1]，表示每个词汇项的概率。每行中的值总和为 1，形成有效的分布。
uniform_sample (relax.Expr) – 均匀采样的 2-D 张量，形状为 (n, 1)。值范围为 0 到 1，表示均匀采样的概率。
sample_indices (relax.Expr) – 形状为 [n, 1] 的 2-D 张量，指示要从中采样的特定概率分布。 sample_indices[i] 的值确定第 i 个 token 应从第 sample_indices[i] 个概率分布中采样。例如，如果有 3 个不同的概率分布，并且要求分别从每个分布中采样 2、3 和 4 个 token，则 sample_indices 将为 [0, 0, 1, 1, 1, 2, 2, 2, 2]。
dtype (str) – 输出张量的数据类型。

返回:

result – 计算出的张量，形状为 (n, 1)。

返回类型:

relax.Expr

示例

prob = [[0.2, 0.3, 0.5], [0.3, 0.4, 0.3]]
usample = [[0.4], [0.9]]
sample_indices = [[0], [1]]

multinomial_from_uniform(prob, usample)
-> [[1], [2]]
multinomial_from_uniform(prob, usample, sample_indices)
-> [[1], [2]]

tvm.relax.op.argmax(x: RelaxExpr, axis: int | None = None, keepdims: bool = False) → RelaxExpr

计算给定轴上张量元素的最大值的索引。

参数:

x (relax.Expr) – 输入数据张量
axis (Optional[int]) – 执行 argmax 操作的轴。默认值 axis=None 将计算输入张量中所有元素的最大值的索引。支持负索引。
keepdims (bool) – 如果设置为 True，则被缩减的轴将保留在结果中，作为大小为 1 的维度。使用此选项，结果将与输入张量正确广播。

返回:

result – 计算结果。

返回类型:

relax.Expr

tvm.relax.op.argmin(x: RelaxExpr, axis: int | None = None, keepdims: bool = False) → RelaxExpr

计算给定轴上张量元素的最小值的索引。

参数:

x (relax.Expr) – 输入数据张量
axis (Optional[int]) – 执行 argmin 操作的轴。默认值 axis=None 将计算输入张量中所有元素的最小值的索引。支持负索引。
keepdims (bool) – 如果设置为 True，则被缩减的轴将保留在结果中，作为大小为 1 的维度。使用此选项，结果将与输入张量正确广播。

返回:

result – 计算结果。

返回类型:

relax.Expr

tvm.relax.op.where(condition: RelaxExpr, x1: RelaxExpr, x2: RelaxExpr) → RelaxExpr

根据条件的值从输入张量中选择元素。

对于给定的位置，如果 condition 为 True，则返回 x1 中的相应值，否则返回 x2 中的相应值。

参数:

condition (relax.Expr) – 当为 True 时，产生 x1；否则，产生 x2。必须与 x1 和 x2 广播兼容。必须具有布尔数据类型。
x1 (relax.Expr) – 第一个输入张量。必须与 condition 和 x2 广播兼容。
x2 (relax.Expr) – 第二个输入张量。必须与 condition 和 x1 广播兼容。

返回:

result – 结果张量。

返回类型:

relax.Expr

tvm.relax.op.nonzero(x: RelaxExpr) → RelaxExpr

查找张量中非零元素的索引。

参数:: x (relax.Expr) – 输入数据张量。
返回:: result – 包含非零元素索引的 2-D 张量。
返回类型:: relax.Expr

Note

此函数等效于 onnx.nonzero。

示例

x = [[0, 1],
     [2, 0]]
nonzero(x) = [[0, 1],
              [1, 0]]

查找给定张量中的唯一元素。此外，它还可以选择性地返回：- 给出唯一值的输入张量的索引；- 重构输入张量的唯一张量的索引；- 每个唯一值在输入张量中出现的次数。

参数:

x (relax.Expr) – 输入张量。
sorted (Union[bool, Expr]) – 是否在作为输出返回之前按升序对唯一元素进行排序。
return_index (Union[bool, Expr]) – 是否返回一个额外的张量，其中包含唯一张量中的元素来自原始输入的索引。
return_inverse (Union[bool, Expr]) – 是否返回一个额外的张量，其中包含原始输入中的元素在返回的唯一列表中最终位置的索引。
return_counts (Union[bool, Expr]) – 是否返回一个额外的张量，其中包含每个唯一元素的计数。
axis (Optional) – 应用 unique 的维度。如果未指定，则返回展平输入的唯一值。

返回:

ret – 使用以下内容创建的 relax 调用

返回类型:

relax.Expr

tvm.relax.op.argsort(data: RelaxExpr, axis: int = -1, descending: bool = False, dtype: str = 'int32')

沿给定轴执行排序，并返回一个索引数组，该数组的形状与输入数组相同，用于按排序顺序索引数据。

参数:

data (relax.Expr) – 输入数据张量。
axis (int) – 对输入张量进行排序的轴。
descending (bool) – 是否按降序排序，默认为 False
dtype (str) – 输出索引的数据类型。

返回:

out – 形状与数据相同的张量。

返回类型:

relax.Expr

tvm.relax.op.sort(x: RelaxExpr, axis: int = -1, descending: bool = False)

沿给定轴执行排序，并返回排序后的数组。

参数:

x (relax.Expr) – 输入张量。
axis (int) – 对输入张量进行排序的轴。默认情况下，使用输入的最后一个轴。
descending (bool) – 是否按降序排序，默认为 False

返回:

out – 排序后的张量。

返回类型:

relax.Expr

tvm.relax.op.topk(data: RelaxExpr, k: int = 1, axis: int = -1, ret_type: str = 'both', largest: bool = True, dtype: str = 'int32')

获取输入张量沿给定轴的前 k 个元素。

ret_type 指定返回类型，可以是 (“both”, “values”, “indices”) 之一。

参数:

data (relax.Expr) – 输入数据张量。
k (int) – 要选择的顶部元素数。如果 k < 1，则返回所有元素。
axis (int) – 对输入张量进行排序的轴。
ret_type (str) – 返回类型 [both, values, indices]。 “both”：同时返回前 k 个数据和索引。“values”：仅返回前 k 个数据。“indices”：仅返回前 k 个索引。
largest (bool) – 是否返回最大或最小元素。如果 largest 为 False，则返回 k 个最小元素。
dtype (str) – 索引输出的数据类型。

返回:

out – 计算结果。

返回类型:

relax.Expr 或 List[relax.Expr]

tvm.relax.op.cumprod(data: RelaxExpr, axis: int | None = None, dtype: str | DataType | None = None, exclusive: bool = False)

Numpy 风格的 cumprod 运算。返回给定轴上元素的累积乘积。

参数:

data (relax.Expr) – 运算符的输入数据。
axis (Optional[int]) – 计算累积乘积的轴。默认值 (None) 是计算展平数组的 cumprod。
dtype (Optional[Union[str, DataType]]) – 返回数组以及在其中计算元素的累加器类型。如果未指定 dtype，则默认为 data 的 dtype。
exclusive (bool) – 如果为 false（默认值），则所有元素都包含在乘积中。如果为 true，则第一个元素将从乘积中排除。

返回:

result – 结果的大小与 data 相同，如果 axis 不为 None，则形状与 data 相同。如果 axis 为 None，则结果为 1-d 数组。

返回类型:

relax.Expr

示例

a = [[1, 2, 3], [4, 5, 6]]

cumprod(a)  # if axis is not provided, cumprod is done over the flattened input.
-> [ 1,  2,  6, 24, 120, 720]

cumprod(a, dtype="float32")
-> [  1.,  2.,  6., 24., 120., 720.]

cumprod(a, axis=0)  # multiply over rows for each of the 3 columns
-> [[1, 2, 3],
    [4, 10, 18]]

cumprod(a, axis=1)
-> [[ 1,  2,  6],
    [ 4,  20, 120]]

a = [1, 1, 1, 0, 1, 1, 0]  # a is a boolean array
cumprod(a, dtype=int32)  # dtype should be provided to get the expected results
-> [1, 1, 1, 0, 0, 0, 0]

tvm.relax.op.cumsum(data: RelaxExpr, axis: int | None = None, dtype: str | DataType | None = None, exclusive: bool = False)

Numpy 风格的 cumsum 运算。返回给定轴上元素的累积包含和。

参数:

data (relax.Expr) – 运算符的输入数据。
axis (Optional[int]) – 计算累积和的轴。默认值 (None) 是计算展平数组的 cumsum。
dtype (Optional[Union[str, DataType]]) – 返回数组以及在其中对元素求和的累加器类型。如果未指定 dtype，则默认为 data 的 dtype。
exclusive (bool) – 如果为 false（默认值），则所有元素都包含在总和中。如果为 true，则第一个元素将从总和中排除。

返回:

result – 结果的大小与 data 相同，如果 axis 不为 None，则形状与 data 相同。如果 axis 为 None，则结果为 1-d 数组。

返回类型:

relax.Expr

示例

a = [[1, 2, 3], [4, 5, 6]]

cumsum(a)  # if axis is not provided, cumsum is done over the flattened input.
-> [ 1,  3,  6, 10, 15, 21]

cumsum(a, dtype="float32")
-> [  1.,   3.,   6.,  10.,  15.,  21.]

cumsum(a, axis=0)  # sum over rows for each of the 3 columns
-> [[1, 2, 3],
    [5, 7, 9]]

cumsum(a, axis=1)
-> [[ 1,  3,  6],
    [ 4,  9, 15]]

a = [1, 0, 1, 0, 1, 1, 0]  # a is a boolean array
cumsum(a, dtype=int32)  # dtype should be provided to get the expected results
-> [1, 1, 2, 2, 3, 4, 4]

tvm.relax.op.max(x: RelaxExpr, axis: int | List[int] | None = None, keepdims: bool = False) → RelaxExpr

计算给定轴上张量元素的最大值。

参数:

x (relax.Expr) – 输入数据张量
axis (Optional[Union[int, List[int]]]) – 执行 max 操作的轴或轴列表。默认值 axis=None 将计算输入张量中所有元素的最大值。支持负索引。
keepdims (bool) – 如果设置为 True，则被缩减的轴将保留在结果中，作为大小为 1 的维度。使用此选项，结果将与输入张量正确广播。

返回:

result – 计算结果。

返回类型:

relax.Expr

tvm.relax.op.mean(x: RelaxExpr, axis: int | List[int] | None = None, keepdims: bool = False) → RelaxExpr

计算给定轴上张量元素的平均值。

参数:

x (relax.Expr) – 输入数据张量
axis (Optional[Union[int, List[int]]]) – 执行 mean 操作的轴或轴列表。默认值 axis=None 将计算输入张量中所有元素的平均值。支持负索引。
keepdims (bool) – 如果设置为 True，则被缩减的轴将保留在结果中，作为大小为 1 的维度。使用此选项，结果将与输入张量正确广播。

返回:

result – 计算结果。

返回类型:

relax.Expr

tvm.relax.op.min(x: RelaxExpr, axis: int | List[int] | None = None, keepdims: bool = False) → RelaxExpr

计算给定轴上张量元素的最小值。

参数:

x (relax.Expr) – 输入数据张量
axis (Optional[Union[int, List[int]]]) – 执行 min 操作的轴或轴列表（可选）。默认值 axis=None 将计算输入张量中所有元素的最小值。支持负索引。
keepdims (bool) – 如果设置为 True，则被缩减的轴将保留在结果中，作为大小为 1 的维度。使用此选项，结果将与输入张量正确广播。

返回:

result – 计算结果。

返回类型:

relax.Expr

tvm.relax.op.prod(x: RelaxExpr, axis: int | List[int] | None = None, keepdims: bool = False) → RelaxExpr

计算给定轴上张量元素的乘积。

参数:

x (relax.Expr) – 输入数据张量
axis (Optional[Union[int, List[int]]]) – 执行乘积运算的轴或轴列表（可选）。默认值 axis=None 将计算输入张量中所有元素的乘积。支持负索引。
keepdims (bool) – 如果设置为 True，则被缩减的轴将保留在结果中，作为大小为 1 的维度。使用此选项，结果将与输入张量正确广播。

返回:

result – 计算结果。

返回类型:

relax.Expr

tvm.relax.op.std(x: RelaxExpr, axis: int | List[int] | None = None, keepdims: bool = False) → RelaxExpr

计算给定轴上张量元素的标准差。

参数:

x (relax.Expr) – 输入数据张量
axis (Optional[Union[int, List[int]]]) – 执行标准差运算的轴或轴列表（可选）。默认值 axis=None 将计算输入张量中所有元素的标准差。支持负索引。
keepdims (bool) – 如果设置为 True，则被缩减的轴将保留在结果中，作为大小为 1 的维度。使用此选项，结果将与输入张量正确广播。

返回:

result – 计算结果。

返回类型:

relax.Expr

tvm.relax.op.sum(x: RelaxExpr, axis: int | List[int] | None = None, keepdims: bool = False) → RelaxExpr

计算给定轴上张量元素的和。

参数:

x (relax.Expr) – 输入数据张量
axis (Optional[Union[int, List[int]]]) – 执行求和运算的轴或轴列表（可选）。默认值 axis=None 将对输入张量中的所有元素求和。支持负索引。
keepdims (bool) – 如果设置为 True，则被缩减的轴将保留在结果中，作为大小为 1 的维度。使用此选项，结果将与输入张量正确广播。

返回:

result – 计算结果。

返回类型:

relax.Expr

tvm.relax.op.variance(x: RelaxExpr, axis: int | List[int] | None = None, keepdims: bool = False) → RelaxExpr

计算给定轴上张量元素的方差。

参数:

x (relax.Expr) – 输入数据张量
axis (Optional[Union[int, List[int]]]) – 执行方差运算的轴或轴列表（可选）。默认值 axis=None 将计算输入张量中所有元素的方差。支持负索引。
keepdims (bool) – 如果设置为 True，则被缩减的轴将保留在结果中，作为大小为 1 的维度。使用此选项，结果将与输入张量正确广播。

返回:

result – 计算结果。

返回类型:

relax.Expr

tvm.relax.op.ewise_fma(x1: RelaxExpr, x2: RelaxExpr, x3: RelaxExpr) → RelaxExpr

逐元素融合乘加运算符。返回 \(x1 * x2 + x3\) 的逐元素结果

参数:

x1 (relax.Expr) – 乘法的左操作数
x2 (relax.Expr) – 乘法的右操作数
x3 (relax.Expr) – 加法的操作数

返回:

result – 计算结果。

返回类型:

relax.Expr

tvm.relax.op.abs(x: RelaxExpr) → RelaxExpr

计算输入数据的逐元素绝对值。

参数:: x (relax.Expr) – 输入数据
返回:: result – 计算结果。
返回类型:: relax.Expr

tvm.relax.op.acos(x: RelaxExpr) → RelaxExpr

计算输入数据的逐元素反余弦值。

参数:: x (relax.Expr) – 输入数据
返回:: result – 计算结果。
返回类型:: relax.Expr

Note

输入张量需要具有浮点 dtype

tvm.relax.op.acosh(x: RelaxExpr) → RelaxExpr

计算输入数据的逐元素反双曲余弦值。

参数:: x (relax.Expr) – 输入数据
返回:: result – 计算结果。
返回类型:: relax.Expr

Note

输入张量需要具有浮点 dtype

tvm.relax.op.asin(x: RelaxExpr) → RelaxExpr

计算输入数据的逐元素反正弦值。

参数:: x (relax.Expr) – 输入数据
返回:: result – 计算结果。
返回类型:: relax.Expr

Note

输入张量需要具有浮点 dtype

tvm.relax.op.asinh(x: RelaxExpr) → RelaxExpr

计算输入数据的逐元素反双曲正弦值。

参数:: x (relax.Expr) – 输入数据
返回:: result – 计算结果。
返回类型:: relax.Expr

Note

输入张量需要具有浮点 dtype

tvm.relax.op.atan(x: RelaxExpr) → RelaxExpr

计算输入数据的逐元素反正切值。

参数:: x (relax.Expr) – 输入数据
返回:: result – 计算结果。
返回类型:: relax.Expr

Note

输入张量需要具有浮点 dtype

tvm.relax.op.atanh(x: RelaxExpr) → RelaxExpr

计算输入数据的逐元素反双曲正切值。

参数:: x (relax.Expr) – 输入数据
返回:: result – 计算结果。
返回类型:: relax.Expr

Note

输入张量需要具有浮点 dtype

tvm.relax.op.bitwise_not(x: RelaxExpr) → RelaxExpr

计算输入数据的按位 NOT。

参数:: x (relax.Expr) – 输入数据
返回:: result – 计算结果。
返回类型:: relax.Expr

tvm.relax.op.ceil(x: RelaxExpr) → RelaxExpr

对输入数据向上取整。

参数:: x (relax.Expr) – 输入数据
返回:: result – 计算结果。
返回类型:: relax.Expr

tvm.relax.op.clip(x: RelaxExpr, min: RelaxExpr, max: RelaxExpr) → RelaxExpr

将张量值裁剪到指定的最小值和最大值之间。

参数:

x (relax.Expr) – 输入数据
min (relax.Expr) – 最小值
max (relax.Expr) – 最大值

返回:

result – 计算结果。

返回类型:

relax.Expr

tvm.relax.op.cos(x: RelaxExpr) → RelaxExpr

计算输入数据的逐元素余弦值。

参数:: x (relax.Expr) – 输入数据
返回:: result – 计算结果。
返回类型:: relax.Expr

Note

输入张量需要具有浮点 dtype

tvm.relax.op.cosh(x: RelaxExpr) → RelaxExpr

计算输入数据的逐元素双曲余弦值。

参数:: x (relax.Expr) – 输入数据
返回:: result – 计算结果。
返回类型:: relax.Expr

Note

输入张量需要具有浮点 dtype

tvm.relax.op.erf(x: RelaxExpr) → RelaxExpr

计算输入的误差函数。

参数:: x (relax.Expr) – 输入数据
返回:: result – 每个元素的计算误差函数。
返回类型:: relax.Expr

tvm.relax.op.exp(x: RelaxExpr) → RelaxExpr

计算数据的逐元素指数。

参数:: x (relax.Expr) – 输入数据
返回:: result – 计算结果。
返回类型:: relax.Expr

Note

输入张量需要具有浮点 dtype

tvm.relax.op.floor(x: RelaxExpr) → RelaxExpr

对输入数据向下取整。

参数:: x (relax.Expr) – 输入数据
返回:: result – 计算结果。
返回类型:: relax.Expr

tvm.relax.op.isfinite(x: RelaxExpr) → RelaxExpr

检查输入值是否为有限值。

参数:: x (relax.Expr) – 输入数据
返回:: result – 计算结果。
返回类型:: relax.Expr

tvm.relax.op.isinf(x: RelaxExpr) → RelaxExpr

检查输入值是否为无穷大。

参数:: x (relax.Expr) – 输入数据
返回:: result – 计算结果。
返回类型:: relax.Expr

tvm.relax.op.isnan(x: RelaxExpr) → RelaxExpr

检查输入值是否为 Nan。

参数:: x (relax.Expr) – 输入数据
返回:: result – 计算结果。
返回类型:: relax.Expr

tvm.relax.op.log(x: RelaxExpr) → RelaxExpr

计算输入数据的逐元素自然对数。

参数:: x (relax.Expr) – 输入数据
返回:: result – 计算结果。
返回类型:: relax.Expr

Note

输入张量需要具有浮点 dtype

tvm.relax.op.logical_not(x: RelaxExpr) → RelaxExpr

计算输入数据的逻辑 NOT。

参数:: x (relax.Expr) – 输入数据
返回:: result – 计算结果。
返回类型:: relax.Expr

tvm.relax.op.negative(x: RelaxExpr) → RelaxExpr

计算输入数据的逐元素负值。

参数:: x (relax.Expr) – 输入数据
返回:: result – 计算结果
返回类型:: relax.Expr

tvm.relax.op.round(x: RelaxExpr) → RelaxExpr

将输入数据的每个元素四舍五入到最接近的整数。

参数:: x (relax.Expr) – 输入数据
返回:: result – 计算结果。
返回类型:: relax.Expr

tvm.relax.op.rsqrt(x: RelaxExpr) → RelaxExpr

计算输入数据的逐元素倒数平方根。

\[1/sqrt(x)\]

参数:: x (relax.Expr) – 输入数据
返回:: result – 计算结果。
返回类型:: relax.Expr

Note

输入张量需要具有浮点 dtype

tvm.relax.op.sigmoid(x: RelaxExpr) → RelaxExpr

计算输入数据的逐元素 Sigmoid 值。

参数:: x (relax.Expr) – 输入数据
返回:: result – 计算结果。
返回类型:: relax.Expr

Note

输入张量需要具有浮点 dtype

tvm.relax.op.sign(x: RelaxExpr) → RelaxExpr

返回输入数据每个元素的符号指示。

参数:: x (relax.Expr) – 输入数据
返回:: result – 计算结果。
返回类型:: relax.Expr

tvm.relax.op.sin(x: RelaxExpr) → RelaxExpr

计算输入数据的逐元素正弦值。

参数:: x (relax.Expr) – 输入数据
返回:: result – 计算结果。
返回类型:: relax.Expr

Note

输入张量需要具有浮点 dtype

tvm.relax.op.sinh(x: RelaxExpr) → RelaxExpr

计算输入数据的逐元素双曲正弦值。

参数:: x (relax.Expr) – 输入数据
返回:: result – 计算结果。
返回类型:: relax.Expr

Note

输入张量需要具有浮点 dtype

tvm.relax.op.sqrt(x: RelaxExpr) → RelaxExpr

计算输入数据的逐元素平方根。

参数:: x (relax.Expr) – 输入数据
返回:: result – 计算结果。
返回类型:: relax.Expr

Note

输入张量需要具有浮点 dtype

tvm.relax.op.square(x: RelaxExpr) → RelaxExpr

计算输入数据每个元素的平方。

参数:: x (relax.Expr) – 输入数据
返回:: result – 计算结果。
返回类型:: relax.Expr

tvm.relax.op.tan(x: RelaxExpr) → RelaxExpr

计算输入数据的逐元素正切值。

参数:: x (relax.Expr) – 输入数据
返回:: result – 计算结果。
返回类型:: relax.Expr

Note

输入张量需要具有浮点 dtype

tvm.relax.op.tanh(x: RelaxExpr) → RelaxExpr

计算输入数据的逐元素双曲正切值。

参数:: x (relax.Expr) – 输入数据
返回:: result – 计算结果。
返回类型:: relax.Expr

Note

输入张量需要具有浮点 dtype

tvm.relax.op.nn

神经网络相关运算符。

tvm.relax.op.nn.adaptive_avg_pool1d(data: RelaxExpr, output_size: int | Tuple[int] | None = None, layout: str = 'NCW', out_layout: str | None = None) → RelaxExpr

1D 自适应平均池化运算符。此运算符为实验性。

此运算符将数据作为输入，并对由 W 表示的每个窗口执行 1D 平均值计算。

在默认情况下，当 data_layout 为 NCW 时，输入数据张量的形状为 (batch_size, in_channels, width)，以生成形状为 (batch_size, in_channels, output_width) 的输出张量。

池化核和步幅大小会自动选择以达到所需的输出大小。

对于 output_size

如果未提供此参数，则输入的高度和宽度将用作输出宽度。

如果为 output_size 提供单个整数，则对于任何输入 (NCW)，输出大小为 (N x C x output_size)。

参数:

data (relax.Expr) – 运算符的输入数据。
output_size (Optional[Union[int, Tuple[int, int]]]) – 输出的高度和宽度。如果未指定，则与输入的高度和宽度相同。如果指定，则长度必须为 1 或 2。
layout (str) – 输入的布局（Layout）。
out_layout (Optional[str]) – 输出的布局（Layout）。如果未指定，则与 data_layout 相同

返回:

result – 计算结果。

返回类型:

relax.Expr

tvm.relax.op.nn.adaptive_avg_pool2d(data: RelaxExpr, output_size: int | Tuple[int, int] | None = None, layout: str = 'NCHW', out_layout: str | None = None) → RelaxExpr

2D 自适应平均池化算子。此算子为实验性功能。

此算子将数据作为输入，并在由 WxH 表示的每个窗口上执行 2D 平均值计算。

在默认情况下，当 data_layout 为 NCHW 时，输入数据张量的形状为 (batch_size, in_channels, height, width)，输出张量的形状为 (batch_size, in_channels, output_height, output_width)。

池化核和步幅大小会自动选择以达到所需的输出大小。

对于 output_size

如果未提供此参数，则输入的高度和宽度将用作输出的高度和宽度。

如果为 output_size 提供单个整数，则对于任何输入 (NCHW)，输出大小为 (N x C x output_size x output_size)。

如果为 output_size 提供整数元组 (height, width)，则对于任何输入 (NCHW)，输出大小为 (N x C x height x width)。

参数:

data (relax.Expr) – 运算符的输入数据。
output_size (Optional[Union[int, Tuple[int, int]]]) – 输出的高度和宽度。如果未指定，则与输入的高度和宽度相同。如果指定，则长度必须为 1 或 2。
layout (str) – 输入的布局（Layout）。
out_layout (Optional[str]) – 输出的布局（Layout）。如果未指定，则与 data_layout 相同

返回:

result – 计算结果。

返回类型:

relax.Expr

tvm.relax.op.nn.adaptive_avg_pool3d(data: RelaxExpr, output_size: int | Tuple[int, int] | None = None, layout: str = 'NCDHW', out_layout: str | None = None) → RelaxExpr

3D 自适应平均池化算子。此算子为实验性功能。

此算子将数据作为输入，并在由 WxH 表示的每个窗口上执行 3D 平均值计算。

在默认情况下，当 data_layout 为 NCDHW 时，输入数据张量的形状为 (batch_size, in_channels, depth, height, width)，输出张量的形状为 (batch_size, in_channels, output_depth, output_height, output_width)。

池化核和步幅大小会自动选择以达到所需的输出大小。

对于 output_size

如果未提供此参数，则输入的深度、高度和宽度将用作输出的深度、高度和宽度。

如果为 output_size 提供单个整数，则对于任何输入 (NCDHW)，输出大小为 (N x C x output_size x output_size x output_size)。

如果为 output_size 提供整数元组 (depth, height, width)，则对于任何输入 (NCDHW)，输出大小为 (N x C x depth x height x width)。

参数:

data (relax.Expr) – 运算符的输入数据。
output_size (Optional[Union[int, Tuple[int, int]]]) – 输出的高度和宽度。如果未指定，则与输入的高度和宽度相同。如果指定，则长度必须为 1 或 3。
layout (str) – 输入的布局（Layout）。
out_layout (Optional[str]) – 输出的布局（Layout）。如果未指定，则与 data_layout 相同

返回:

result – 计算结果。

返回类型:

relax.Expr

tvm.relax.op.nn.attention(query: RelaxExpr, key: RelaxExpr, value: RelaxExpr, bias: RelaxExpr | None = None, scale: FloatImm | None = None, causal_mask: str | None = None, window_size: int | None = None) → RelaxExpr

计算融合多头注意力。

所有输入张量均为 BSNH 布局的 4 维张量。

\[FMA(Q, K, V) = \text{Softmax}(Q @ K^T) @ V\]

Note

输入张量需要具有 float16 dtype

参数:

query (relax.Expr) – 算子的输入查询（query）。输入查询的布局应为 (batch_size, seq_len, num_head, head_dim)。
key (relax.Expr) – 算子的输入键（key）。输入键的布局应为 (batch_size, seq_len_kv, num_head, head_dim)。
value (relax.Expr) – 算子的输入值（value）。输入值的布局应为 (batch_size, seq_len_kv, num_head, head_dim_v)。
bias (Optional[Expr]) – 算子的可选注意力偏置（bias）。注意力偏置的布局应为以 seq_len_kv 结尾的 4 维张量，并且可广播到 (batch_size, num_head, seq_len, seq_len_kv)。
scale (Optional[float]) – 应用于注意力分数的缩放值（scale value），默认为 1 / sqrt(head_dim)。
causal_mask (Optional[str]) –
可选的因果掩码（causal mask），即 ‘TopLeft’ 和 ‘BottomRight’。对于 ‘TopLeft’，掩码矩阵类似于 np.tril(*, k=0)，而对于 ‘BottomRight’，掩码矩阵类似于 np.tril(*, k=abs(seq_len - seq_len_kv))。例如，当 seq_len = 4，seq_len_kv = 2 时，‘TopLeft’ 的掩码为
```
[[1, 0],
[1, 1],
[1, 1],
[1, 1]]
```
‘BottomRight’ 的掩码为
```
[[1, 1],
[1, 1],
[1, 1],
[1, 1]]
```
当 seq_len = 2，seq_len_kv = 4 时，‘TopLeft’ 的掩码为
```
[[1, 0, 0, 0],
[1, 1, 0, 0]]
```
‘BottomRight’ 的掩码为
```
[[1, 1, 1, 0],
[1, 1, 1, 1]]
```
window_size (Optional[int]) – 滑动窗口注意力（sliding-window attention）的窗口大小。

返回:

result – 计算结果。输出的布局应为 (batch_size, seq_len, num_head, head_dim_v)。

返回类型:

relax.Expr

tvm.relax.op.nn.attention_var_len(queries: RelaxExpr, keys: RelaxExpr, values: RelaxExpr, seqstart_q: RelaxExpr, max_seqlen_q: RelaxExpr, seqstart_k: RelaxExpr | None = None, max_seqlen_k: RelaxExpr | None = None, scale: FloatImm | None = None, causal_mask: str | None = None, window_size: int | None = None) → RelaxExpr

计算变长批次序列上的融合多头注意力。

给定连接的输入和序列长度信息，此算子比为每个序列单独调用普通注意力算子更有效地计算所有序列的注意力。

参数:

queries (relax.Expr) – 沿第二个轴连接的输入查询。其形状必须为 (1, total_seq_len, num_head, head_dim)。
keys (relax.Expr) – 沿第二个轴连接的输入键。其形状必须为 (1, total_seq_len_kv, num_head, head_dim)。
values (relax.Expr) – 沿第二个轴连接的输入值。其形状必须为 (1, total_seq_len_kv, num_head, head_dim_v)。
seqstart_q (Optional[Expr]) – 查询序列长度的累积和（cumsum），前面补 0。其 dtype 必须为 int32。例如，如果批处理的序列长度为 [2, 5, 3]，则此张量的值为 [0, 2, 7, 10]。
seqstart_k (Optional[Expr]) – 键序列长度的累积和，前面补 0。默认情况下，它与 seqstart_q 相同。
max_seqlen_q (Optional[Expr]) – 批次中最大查询序列长度。它必须是 int32。
max_seqlen_k (Optional[Expr]) – 批次中最大键序列长度。它必须是 int32。默认情况下，它与 max_seqlen_q 相同。
scale (Optional[float]) – 应用于注意力分数的缩放值（scale value），默认为 1 / sqrt(head_dim)。
causal_mask (Optional[str]) –
可选的因果掩码（causal mask），即 ‘TopLeft’ 和 ‘BottomRight’。对于 ‘TopLeft’，掩码矩阵类似于 np.tril(*, k=0)，而对于 ‘BottomRight’，掩码矩阵类似于 np.tril(*, k=abs(seq_len - seq_len_kv))。例如，当 seq_len = 4，seq_len_kv = 2 时，‘TopLeft’ 的掩码为
```
[[1, 0],
[1, 1],
[1, 1],
[1, 1]]
```
‘BottomRight’ 的掩码为
```
[[1, 1],
[1, 1],
[1, 1],
[1, 1]]
```
当 seq_len = 2，seq_len_kv = 4 时，‘TopLeft’ 的掩码为
```
[[1, 0, 0, 0],
[1, 1, 0, 0]]
```
‘BottomRight’ 的掩码为
```
[[1, 1, 1, 0],
[1, 1, 1, 1]]
```
window_size (Optional[int]) – 滑动窗口注意力（sliding-window attention）的窗口大小。

返回:

result – 计算结果，形状为 (1, total_seq_len, num_head, head_dim_v)。

返回类型:

relax.Expr

tvm.relax.op.nn.avg_pool1d(data: RelaxExpr, pool_size: int | Tuple[int, int] = (1,), strides: int | Tuple[int, int] = (1,), padding: int | Tuple[int, ...] = (0, 0), dilation: int | Tuple[int, int] = (1,), ceil_mode: bool = False, count_include_pad: bool = False, layout: str = 'NCW', out_layout: str | None = None) → RelaxExpr

1D 平均池化算子。

此算子将数据作为输入，并在 pool_size 大小的窗口中，以 stride 定义的步长进行 1D 平均值计算。

在默认情况下，当 data_layout 为 NCW 时，输入数据张量的形状为 (batch_size, channels, width)，以生成输出张量。

ceil_mode 用于在计算输出形状时选择向上取整或向下取整。count_include_pad 指示在计算中是否包括或排除填充的输入值。此算子接受数据布局规范。

参数:

data (relax.Expr) – 运算符的输入数据。
pool_size (Union[int, Tuple[int]]) – 池化窗口的大小。长度必须为 1。
strides (Union[int, Tuple[int]]) – 池化的步长。长度必须为 1。
padding (Union[int, Tuple[int, int]]) – 池化的填充（padding）。长度必须为 1 或 2。
dilation (Union[int, Tuple[int]]) – 池化的空洞（dilation）。长度必须为 1。
ceil_mode (bool) – 一个布尔值，指示是使用向上取整还是向下取整来计算输出形状。通过使用向上取整，输入张量中的每个元素都将被滑动窗口覆盖。
count_include_pad (bool, optional) – 是否包含填充以计算平均值。
layout (str) – 输入的布局（Layout）。
out_layout (Optional[str]) – 输出的布局（Layout）。如果未指定，则与 data_layout 相同

返回:

result – 计算结果。

返回类型:

Expr

tvm.relax.op.nn.avg_pool2d(data: RelaxExpr, pool_size: int | Tuple[int, int] = (1, 1), strides: int | Tuple[int, int] = (1, 1), padding: int | Tuple[int, ...] = (0, 0), dilation: int | Tuple[int, int] = (1, 1), ceil_mode: bool = False, count_include_pad: bool = False, layout: str = 'NCHW', out_layout: str | None = None) → RelaxExpr

2D 平均池化算子。

此算子将数据作为输入，并在 pool_size 大小的窗口中，以 stride 定义的步长进行 2D 平均值计算。

在默认情况下，当 data_layout 为 NCHW 时，输入数据张量的形状为 (batch_size, in_channels, height, width)，以生成符合以下规则的输出张量

对于形状为 (b, c, h, w) 的数据和 pool_size (kh, kw)

\[\mbox{out}(b, c, y, x) = \frac{1}{kh * kw} \sum_{m=0, \ldots, kh-1} \sum_{n=0, \ldots, kw-1} \mbox{data}(b, c, \mbox{stride}[0] * y + m, \mbox{stride}[1] * x + n)\]

填充在计算之前应用于数据。ceil_mode 用于在计算输出形状时选择向上取整或向下取整。此算子接受数据布局规范。

参数:

data (relax.Expr) – 运算符的输入数据。
pool_size (Union[int, Tuple[int, int]]) – 池化窗口的大小。长度必须为 1 或 2。
strides (Union[int, Tuple[int, int]]) – 池化的步长。长度必须为 1 或 2。
padding (Union[int, Tuple[int, ...]]) – 池化的填充（padding）。长度必须为 1、2 或 4。
空洞率（dilation） (Union[int, Tuple[int, int]]) – 池化的空洞率。长度必须为 1 或 2。
ceil_mode (bool) – 一个布尔值，指示是使用向上取整还是向下取整来计算输出形状。通过使用向上取整，输入张量中的每个元素都将被滑动窗口覆盖。
count_include_pad (bool, optional) – 是否包含填充以计算平均值。
layout (str) – 输入的布局（Layout）。
out_layout (Optional[str]) – 输出的布局（Layout）。如果未指定，则与 data_layout 相同

返回:

result – 计算结果。

返回类型:

Expr

tvm.relax.op.nn.avg_pool3d(data: RelaxExpr, pool_size: int | Tuple[int, int] = (1, 1, 1), strides: int | Tuple[int, int] = (1, 1, 1), padding: int | Tuple[int, ...] = (0, 0, 0), dilation: int | Tuple[int, int] = (1, 1, 1), ceil_mode: bool = False, count_include_pad: bool = False, layout: str = 'NCDHW', out_layout: str | None = None) → RelaxExpr

2D 平均池化算子。

此运算符将数据作为输入，并通过步幅定义的 pool_size 大小的窗口执行 3D 平均值计算。

在默认情况下，当 data_layout 为 NCDHW 时，输入数据张量的形状为 (batch_size, channels, depth, height, width)，以生成输出张量。

ceil_mode 用于在计算输出形状时选择向上取整或向下取整。count_include_pad 指示在计算中是否包括或排除填充的输入值。此算子接受数据布局规范。

参数:

data (relax.Expr) – 运算符的输入数据。
pool_size (Union[int, Tuple[int, int, int]]) – 池化窗口的大小。长度必须为 1 或 3。
strides (Union[int, Tuple[int, int, int]]) – 池化的步幅。长度必须为 1 或 3。
padding (Union[int, Tuple[int, ...]]) – 池化的填充。长度必须为 1、3 或 6。
dilation (Union[int, Tuple[int, int, int]]) – 池化的空洞率。长度必须为 1 或 3。
ceil_mode (bool) – 一个布尔值，指示是使用向上取整还是向下取整来计算输出形状。通过使用向上取整，输入张量中的每个元素都将被滑动窗口覆盖。
count_include_pad (bool, optional) – 是否包含填充以计算平均值。
layout (str) – 输入的布局（Layout）。
out_layout (Optional[str]) – 输出的布局（Layout）。如果未指定，则与 data_layout 相同

返回:

result – 计算结果。

返回类型:

Expr

tvm.relax.op.nn.batch_norm(data: RelaxExpr, gamma: RelaxExpr, beta: RelaxExpr, moving_mean: RelaxExpr, moving_var: RelaxExpr, axis: int, epsilon: float = 1e-05, center: bool = True, scale: bool = True, momentum: float = 0.1) → RelaxExpr

批归一化层 (Ioffe and Szegedy, 2014)。

在每个批次中对输入进行归一化，即应用一种变换，使平均激活值接近 0，激活值标准差接近 1。

\[\begin{split}data\_mean[i] = mean(data[:,i,:,...]) \\ data\_var[i] = var(data[:,i,:,...])\end{split}\]

mean 和 var 都将输入视为向量，并返回标量。

然后计算归一化后的输出，其形状与输入相同，如下所示

\[out[:,i,:,...] = \frac{data[:,i,:,...] - data\_mean[i]}{\sqrt{data\_var[i]+\epsilon}} * gamma[i] + beta[i]\]

假设输入在轴 1 上的大小为 k，则 gamma 和 beta 的形状均为 (k,)。

除了输入和输出之外，此运算符还接受两个辅助状态，moving_mean 和 moving_var，它们是长度为 k 的向量。它们是整个数据集的全局统计信息，通过以下方式更新：

moving_mean = moving_mean * momentum + data_mean * (1 - momentum)
moving_var = moving_var * momentum + data_var * (1 - momentum)

参数 axis 指定输入形状的哪个轴表示“通道”（单独归一化的组）。默认值为 1。指定 -1 会将通道轴设置为输入形状中的最后一项。

Note

此运算符有两种模式

训练模式。
- 使用从**此**批次计算的均值和方差进行归一化。
- 更新并返回运行均值和运行方差。
推理模式。
- 使用 running_mean 和 running_var 参数进行归一化。
- 不更新运行均值和运行方差。仅返回原始值。

在合法化阶段，此运算符默认将被合法化为训练模式。

您可以使用 tvm.relax.transform.DecomposeOpsForInference 来分解运算符，使其执行推理模式计算。类似地，使用 tvm.relax.transform.DecomposeOpsForTraining 来执行训练模式计算。

参数:

data (relax.Expr) – 运算符的输入数据。
gamma (relax.Expr) – Gamma 缩放因子。
beta (relax.Expr) – Beta 偏移因子。
moving_mean (relax.Expr) – 输入的运行均值。
moving_var (relax.Expr) – 输入的运行方差。
axis (int) – 应用归一化的轴。
epsilon (float) – 添加到方差中的小浮点数，以避免除以零。
center (bool) – 指示是否将 beta 偏移添加到归一化张量。
scale (bool) – 指示是否将乘以 gamma 缩放。
momentum (float) – 用于 moving_mean 和 moving_var 更新的值。

返回:

result – 计算结果。

返回类型:

relax.Expr

tvm.relax.op.nn.conv1d(data: RelaxExpr, weight: RelaxExpr, strides: int | Tuple[int] = 1, padding: int | Tuple[int, ...] = 0, dilation: int | Tuple[int] = 1, groups: int = 1, data_layout: str = 'NCW', kernel_layout: str = 'OIW', out_layout: str | None = None, out_dtype: str | DataType | None = None) → RelaxExpr

1D 卷积。

此运算符将权重作为 1D 卷积核，并将其与数据进行卷积以产生输出。

在默认情况下，当 data_layout 为 NCW 且 kernel_layout 为 OIW 时，conv1d 接受形状为 (batch_size, in_channels, width) 的数据张量和形状为 (channels, in_channels, kernel_w) 的权重张量，其中 kernel_w 是 W 核维度的长度，以生成具有以下规则的输出张量

\[\mbox{out}[b, c, x] = \sum_{dx, k} \mbox{data}[b, k, \mbox{strides} * x + dx] * \mbox{weight}[c, k, dx]\]

填充和空洞分别应用于数据和权重，然后再进行计算。此运算符接受数据布局规范。从语义上讲，该运算符会将布局转换为规范布局（数据为 NCW，权重为 OIW），执行计算，然后转换为 out_layout。

参数:

data (relax.Expr) – 运算符的输入数据。
weight (relax.Expr) – 权重表达式。
strides (Union[int, Tuple[int]]) – 卷积的步幅。长度必须为 1。
padding (Union[int, Tuple[int, ...]]) – 卷积前在输入两侧的填充。长度必须为 1 或 2。
dilation (Union[int, Tuple[int, int]]) – 指定用于空洞卷积的空洞率。长度必须为 1。
groups (int) – 将输入拆分为组以进行分组卷积的组数。输入和输出通道的数量应可被组数整除。
data_layout (str) – 输入的布局。
kernel_layout (str) – 权重的布局。
out_layout (Optional[str]) – 输出的布局（Layout）。如果未指定，则与 data_layout 相同
out_dtype (Optional[Union[str, DataType]]) – 指定混合精度 conv1d 的输出数据类型。

返回:

result – 计算结果。

返回类型:

relax.Expr

tvm.relax.op.nn.conv1d_transpose(data: RelaxExpr, weight: RelaxExpr, strides: int | Tuple[int] = 1, padding: int | Tuple[int, ...] = 0, output_padding: int | Tuple[int] = 0, dilation: int | Tuple[int] = 1, groups: int = 1, data_layout: str = 'NCW', kernel_layout: str = 'IOW', out_layout: str | None = None, out_dtype: str | DataType | None = None) → RelaxExpr

1D 转置卷积运算符。

此运算符可以看作是 conv1d 的梯度运算符。

输出形状可以在简单情况下解释，即 data_layout == “NCW” 且 kernel_layout == “IOW”。假设 data 的形状为 (N, in_channel, in_w)，weight 的形状为 (in_channel, out_channel, weight_w)，我们需要确保 in_channel % groups == 0。输出的形状将为 (N, out_channel * groups, out_w)，其中

out_w = ((in_w - 1) * strides[0] + weight_w - 2 * padding[0] + output_padding[0])

参数:

data (relax.Expr) – 运算符的输入数据。
weight (relax.Expr) – 权重表达式。
strides (Union[int, Tuple[int]]) – 卷积的步幅。长度必须为 1。
padding (Union[int, Tuple[int, ...]]) – 卷积前在输入两侧的填充。长度必须为 1 或 2。
output_padding (Union[int, Tuple[int, ...]], 可选) – 用于消除输出形状的歧义。
dilation (Union[int, Tuple[int]]) – 指定用于空洞卷积的空洞率。长度必须为 1。
groups (int) – 将输入拆分为组以进行分组卷积的组数。输入和输出通道的数量应可被组数整除。
data_layout (str) – 输入的布局。
kernel_layout (str) – 权重的布局。
out_layout (Optional[str]) – 输出的布局（Layout）。如果未指定，则与 data_layout 相同
out_dtype (Optional[Union[str, DataType]]) – 指定混合精度 conv2d 的输出数据类型。

返回:

result – 计算结果。

返回类型:

relax.Expr

tvm.relax.op.nn.conv2d(data: RelaxExpr, weight: RelaxExpr, strides: int | Tuple[int, int] = (1, 1), padding: int | Tuple[int, ...] = (0, 0), dilation: int | Tuple[int, int] = (1, 1), groups: int = 1, data_layout: str = 'NCHW', kernel_layout: str = 'OIHW', out_layout: str | None = None, out_dtype: str | DataType | None = None) → RelaxExpr

二维卷积。

此运算符将权重作为卷积核，并将其与数据进行卷积以产生输出。

在默认情况下，当 data_layout 为 NCHW 且 kernel_layout 为 OIHW 时，conv2d 接受形状为 (batch_size, in_channels, height, width) 的数据张量和形状为 (channels, in_channels, kernel_h, kernel_w) 的权重张量，其中 kernel_h 和 kernel_w 是 H 和 W 核维度的长度，以生成具有以下规则的输出张量

\[\mbox{out}[b, c, y, x] = \sum_{dy, dx, k} \mbox{data}[b, k, \mbox{strides}[0] * y + dy, \mbox{strides}[1] * x + dx] * \mbox{weight}[c, k, dy, dx]\]

在计算之前，填充和空洞分别应用于数据和权重。此运算符接受数据布局规范。从语义上讲，该运算符会将布局转换为规范布局（数据为 NCHW，权重为 OIHW），执行计算，然后转换为 out_layout。

参数:

data (relax.Expr) – 运算符的输入数据。
weight (relax.Expr) – 权重表达式。
strides (Union[int, Tuple[int, int]]) – 卷积的步幅。长度必须为 1 或 2。
padding (Union[int, Tuple[int, ...]]) – 卷积在卷积之前在输入两侧的填充。长度必须为 1、2 或 4。
dilation (Union[int, Tuple[int, int]]) – 指定用于空洞卷积的空洞率。长度必须为 1 或 2。
groups (int) – 将输入拆分为组以进行分组卷积的组数。输入和输出通道的数量应可被组数整除。
data_layout (str) – 输入的布局。
kernel_layout (str) – 权重的布局。
out_layout (Optional[str]) – 输出的布局（Layout）。如果未指定，则与 data_layout 相同
out_dtype (Optional[Union[str, DataType]]) – 指定混合精度 conv2d 的输出数据类型。

返回:

result – 计算结果。

返回类型:

relax.Expr

tvm.relax.op.nn.conv2d_transpose(data: RelaxExpr, weight: RelaxExpr, strides: int | Tuple[int, int] = (1, 1), padding: int | Tuple[int, ...] = (0, 0), output_padding: int | Tuple[int, int] = (0, 0), dilation: int | Tuple[int, int] = (1, 1), groups: int = 1, data_layout: str = 'NCHW', kernel_layout: str = 'IOHW', out_layout: str | None = None, out_dtype: str | DataType | None = None) → RelaxExpr

二维转置卷积运算符。

此运算符旨在作为 conv2d 的梯度运算符。也就是说，如果

out = conv2d(data, weight, strides, padding, dilation),

则 w.r.t. 数据的梯度可以按如下方式计算

data_grad = conv2d_transpose(out_grad, weight, strides, padding, output_padding, dilation),

其中 output_padding 是用于确定输出形状的参数。

当 data_layout == “NCHW” 和 kernel_layout == “IOHW” 时，可以在简单情况下解释输出形状。假设 data 的形状为 (N, in_channel, in_h, in_w)，weight 的形状为 (in_channel, out_channel, weight_h, weight_w)，我们需要确保 in_channel % groups == 0。输出的形状将为 (N, out_channel * groups, out_h, out_w)，其中

out_h = ((in_h - 1) * strides[0] + weight_h - 2 * padding[0] + output_padding[0])
out_w = ((in_w - 1) * strides[1] + weight_w - 2 * padding[1] + output_padding[1])

参数:

data (relax.Expr) – 运算符的输入数据。
weight (relax.Expr) – 权重表达式。
strides (Union[int, Tuple[int, int]]) – 卷积的步幅。长度必须为 1 或 2。
padding (Union[int, Tuple[int, ...]]) – 卷积在卷积之前在输入两侧的填充。长度必须为 1、2 或 4。
output_padding (Union[int, Tuple[int, ...]], 可选) – 用于消除输出形状的歧义。
dilation (Union[int, Tuple[int, int]]) – 指定用于空洞卷积的空洞率。长度必须为 1 或 2。
groups (int) – 将输入拆分为组以进行分组卷积的组数。输入和输出通道的数量应可被组数整除。
data_layout (str) – 输入的布局。
kernel_layout (str) – 权重的布局。
out_layout (Optional[str]) – 输出的布局（Layout）。如果未指定，则与 data_layout 相同
out_dtype (Optional[Union[str, DataType]]) – 指定混合精度 conv2d 的输出数据类型。

返回:

result – 计算结果。

返回类型:

relax.Expr

tvm.relax.op.nn.conv3d(data: RelaxExpr, weight: RelaxExpr, strides: int | Tuple[int, int] = (1, 1, 1), padding: int | Tuple[int, ...] = (0, 0, 0), dilation: int | Tuple[int, int] = (1, 1, 1), groups: int = 1, data_layout: str = 'NCDHW', kernel_layout: str = 'OIDHW', out_layout: str | None = None, out_dtype: str | DataType | None = None) → RelaxExpr

三维卷积。

此运算符将权重作为卷积核，并将其与数据进行卷积以产生输出。

在默认情况下，当 data_layout 为 NCDHW 且 kernel_layout 为 OIDHW 时，conv3d 接受形状为 (batch_size, in_channels, depth, height, width) 的数据张量和形状为 (channels, in_channels, kernel_d, kernel_h, kernel_w) 的权重张量，其中 kernel_d、kernel_h 和 kernel_w 是 D、H 和 W 核维度的长度，以生成具有以下规则的输出张量

\[\mbox{out}[b, c, z, y, x] = \sum_{dz, dy, dx, k} \mbox{data}[b, k, \mbox{strides}[0] * z + dz, \mbox{strides}[1] * y + dy, \mbox{strides}[2] * x + dx] * \mbox{weight}[c, k, dz, dy, dx]\]

在计算之前，填充和空洞分别应用于数据和权重。此运算符接受数据布局规范。从语义上讲，该运算符会将布局转换为规范布局（数据为 NCDHW，权重为 OIDHW），执行计算，然后转换为 out_layout。

参数:

data (relax.Expr) – 运算符的输入数据。
weight (relax.Expr) – 权重表达式。
strides (Union[int, Tuple[int, int, int]]) – 卷积的步幅。长度必须为 1 或 3。
padding (Union[int, Tuple[int, ...]]) – 卷积在卷积之前在输入两侧的填充。长度必须为 1、3 或 6。
dilation (Union[int, Tuple[int, int, int]]) – 指定用于空洞卷积的空洞率。长度必须为 1 或 3。
groups (int) – 将输入拆分为组以进行分组卷积的组数。输入和输出通道的数量应可被组数整除。
data_layout (str) – 输入的布局。
kernel_layout (str) – 权重的布局。
out_layout (Optional[str]) – 输出的布局（Layout）。如果未指定，则与 data_layout 相同
out_dtype (Optional[Union[str, DataType]]) – 指定混合精度 conv2d 的输出数据类型。

返回:

result – 计算结果。

返回类型:

relax.Expr

tvm.relax.op.nn.cross_entropy_with_logits(predictions: RelaxExpr, labels: RelaxExpr) → RelaxExpr

预测和标签之间使用 logits 的 CrossEntropy。

预测和标签的形状必须相同。当 ndim >= 2 时，第一维被视为 batch_size N。在这种情况下，计算结果将除以 N 以执行均值缩减。

\[\text{cross\_entropy\_with\_logits}(x_i, y_i) = \frac{\sum_i -x_i \cdot y_i}{N}\]

参数:

predictions (relax.Expr) – 预测值。
labels (relax.Expr) – 标签（ground truth 值）。

返回:

result – 计算结果。

返回类型:

relax.Expr

tvm.relax.op.nn.dropout(data: RelaxExpr, rate: float = 0.5) → RelaxExpr

将 dropout 操作应用于输入张量。

在训练期间，输入的每个元素都以概率 p 设置为零。整个数组按 1/(1-p) 缩放，以保持输入的期望总和不变。

参数:

data (relax.Expr) – 运算符的输入数据。
rate (float) – 元素重置为 0 的概率。

返回:

result – dropout 的结果，它是两个张量的元组。第一个是原始张量，第二个是掩码张量（未 dropout 的元素为 1.0，dropout 的元素为 0.0）

返回类型:

relax.Expr

tvm.relax.op.nn.gelu(data: RelaxExpr) → RelaxExpr

高斯误差线性单元函数

\[\text{GeLU}(x) = 0.5 * x * (1 + \text{erf}(x * 0.5**0.5))\]

其中 \(erf\) 是高斯误差函数。

参数:: data (relax.Expr) – 输入数据
返回:: result – 计算结果。
返回类型:: relax.Expr

Note

输入张量需要具有浮点 dtype

tvm.relax.op.nn.gelu_tanh(data: RelaxExpr) → RelaxExpr

使用 tanh 近似的高斯误差线性单元函数

\[\text{GELU}(x) = 0.5 * x * (1 + \text{Tanh}(\sqrt(2 / \pi) * (x + 0.044715 * x^3)))\]

参数:: data (relax.Expr) – 输入数据
返回:: result – 计算结果。
返回类型:: relax.Expr

Note

输入张量需要具有浮点 dtype

tvm.relax.op.nn.group_norm(data: RelaxExpr, gamma: RelaxExpr, beta: RelaxExpr, num_groups: int, channel_axis: int, axes: int | List[int], epsilon: float = 1e-05, center: bool = True, scale: bool = True) → RelaxExpr

Group normalization (Yuxin Wu et al., 2016)。将分组归一化应用于 n 维输入数组。此运算符接受 n 维输入数组。首先沿通道轴将输入数组分成组。然后将层归一化应用于每个组。

参数:

data (relax.Expr) – 应用 group_norm 的输入。
gamma (relax.Expr) – Gamma 缩放因子。
beta (relax.Expr) – Beta 偏移因子。
num_groups (int) – 将通道分成的组数。
channel_axis (int) – 输入数据中通道轴的索引。
axes (Union[int, List[int]]) – 应用归一化的轴（不包括组轴）
epsilon (float) – 添加到方差中的小浮点数，以避免除以零。
center (bool) – 指示是否将 beta 偏移添加到归一化张量。
scale (bool) – 指示是否将乘以 gamma 缩放。

返回:

result – 计算结果。

返回类型:

relax.Expr

tvm.relax.op.nn.layer_norm(data: RelaxExpr, gamma: RelaxExpr, beta: RelaxExpr, axes: int | List[int], epsilon: float = 1e-05, center: bool = True, scale: bool = True) → RelaxExpr

Layer normalization (Lei Ba et al., 2016)。将层归一化应用于 n 维输入数组。此运算符接受 n 维输入数组，并使用给定的轴对输入进行归一化

\[out = \frac{data - mean(data, axis)}{\sqrt{var(data, axis)+\epsilon}} * gamma + beta\]

与批量归一化不同，均值和方差是沿通道维度计算的。

假设轴 1 上的输入大小为 k，则 gamma 和 beta 的形状均为 (k,)。

Note

此运算符可以在推理时被优化掉。

参数:

data (relax.Expr) – 应用 layer_norm 的输入。
gamma (relax.Expr) – Gamma 缩放因子。
beta (relax.Expr) – Beta 偏移因子。
axes (Union[int, List[int]]) – 应用归一化的轴。
epsilon (float) – 添加到方差中的小浮点数，以避免除以零。
center (bool) – 指示是否将 beta 偏移添加到归一化张量。
scale (bool) – 指示是否将乘以 gamma 缩放。

返回:

result – 计算结果。

返回类型:

relax.Expr

tvm.relax.op.nn.leakyrelu(data: RelaxExpr, alpha: float = 0.01) → RelaxExpr

线性整流单元。

\[text{LeakyReLU, negative_slope}(x) = max(x, 0) + negative_slope * min(x, 0)\]

参数:

data (relax.Expr) – 输入数据
alpha (float) – 控制负斜率的角度，用于负输入。默认值为 0.01

返回:

result – 计算结果。

返回类型:

relax.Expr

tvm.relax.op.nn.log_softmax(data: RelaxExpr, axis: int = -1) → RelaxExpr

计算 log softmax。

\[\text{log\_softmax}(x_i) = \log\left( \frac{\exp(x_i)}{\sum_j \exp(x_j)}\right)\]

Note

此运算符可以在推理时被优化掉。

参数:

data (relax.Expr) – 运算符的输入数据。
axis (int) – 在计算 log softmax 时要进行求和的轴。如果未指定，则默认为输入张量的最后一个轴。支持负索引。

返回:

result – 计算结果。

返回类型:

relax.Expr

tvm.relax.op.nn.max_pool1d(data: RelaxExpr, pool_size: int | Tuple[int, int] = (1,), strides: int | Tuple[int, int] = (1,), padding: int | Tuple[int, ...] = (0, 0), dilation: int | Tuple[int, int] = (1,), ceil_mode: bool = False, count_include_pad: bool = False, layout: str = 'NCW', out_layout: str | None = None) → RelaxExpr

一维最大池化运算符。

此运算符将数据作为输入，并在 pool_size 大小的窗口中进行 1D 最大值计算，步幅由 stride 定义。

在默认情况下，当 data_layout 为 NCW 时，数据张量的形状为 (batch_size, channels, width)，以生成输出张量。

ceil_mode 用于在计算输出形状时选择向上取整或向下取整。count_include_pad 指示在计算中是否包括或排除填充的输入值。此算子接受数据布局规范。

参数:

data (relax.Expr) – 运算符的输入数据。
pool_size (Union[int, Tuple[int, int]]) – 池化窗口的大小。长度必须为 1。
strides (Union[int, Tuple[int, int]]) – 池化的步幅。长度必须为 1。
padding (Union[int, Tuple[int, ...]]) – 池化的填充。长度必须为 1 或 2。
dilation (Union[int, Tuple[int, int]]) – 池化的扩张。长度必须为 1。
ceil_mode (bool) – 一个布尔值，指示是使用向上取整还是向下取整来计算输出形状。通过使用向上取整，输入张量中的每个元素都将被滑动窗口覆盖。
count_include_pad (bool, optional) – 是否包含填充以计算平均值。
layout (str) – 输入的布局（Layout）。
out_layout (Optional[str]) – 输出的布局（Layout）。如果未指定，则与 data_layout 相同

返回:

result – 计算结果。

返回类型:

Expr

tvm.relax.op.nn.max_pool2d(data: RelaxExpr, pool_size: int | Tuple[int, int] = (1, 1), strides: int | Tuple[int, int] = (1, 1), padding: int | Tuple[int, ...] = (0, 0), dilation: int | Tuple[int, int] = (1, 1), ceil_mode: bool = False, count_include_pad: bool = False, layout: str = 'NCHW', out_layout: str | None = None) → RelaxExpr

二维最大池化运算符。

此运算符将数据作为输入，并在 pool_size 大小的窗口中进行 2D 最大值计算，步幅由 stride 定义。

在默认情况下，当 data_layout 为 NCHW 时，输入数据张量的形状为 (batch_size, in_channels, height, width)，以生成符合以下规则的输出张量

对于形状为 (b, c, h, w) 的数据和 pool_size (kh, kw)

\[\mbox{out}(b, c, y, x) = \max_{m=0, \ldots, kh-1} \max_{n=0, \ldots, kw-1} \mbox{data}(b, c, \mbox{stride}[0] * y + m, \mbox{stride}[1] * x + n)\]

填充在计算之前应用于数据。ceil_mode 用于在计算输出形状时选择向上取整或向下取整。此算子接受数据布局规范。

参数:

data (relax.Expr) – 运算符的输入数据。
pool_size (Union[int, Tuple[int, int]]) – 池化窗口的大小。长度必须为 1 或 2。
strides (Union[int, Tuple[int, int]]) – 池化的步长。长度必须为 1 或 2。
padding (Union[int, Tuple[int, ...]]) – 池化的填充（padding）。长度必须为 1、2 或 4。
空洞率（dilation） (Union[int, Tuple[int, int]]) – 池化的空洞率。长度必须为 1 或 2。
ceil_mode (bool) – 一个布尔值，指示是使用向上取整还是向下取整来计算输出形状。通过使用向上取整，输入张量中的每个元素都将被滑动窗口覆盖。
count_include_pad (bool, optional) – 是否包含填充以计算平均值。
layout (str) – 输入的布局（Layout）。
out_layout (Optional[str]) – 输出的布局（Layout）。如果未指定，则与 data_layout 相同

返回:

result – 计算结果。

返回类型:

Expr

tvm.relax.op.nn.max_pool3d(data: RelaxExpr, pool_size: int | Tuple[int, int] = (1, 1, 1), strides: int | Tuple[int, int] = (1, 1, 1), padding: int | Tuple[int, ...] = (0, 0, 0), dilation: int | Tuple[int, int] = (1, 1, 1), ceil_mode: bool = False, count_include_pad: bool = False, layout: str = 'NCDHW', out_layout: str | None = None) → RelaxExpr

三维最大池化运算符。

此运算符将数据作为输入，并在 pool_size 大小的窗口中进行 3D 最大值计算，步幅由 stride 定义。

在默认情况下，当 data_layout 为 NCDHW 时，输入数据张量的形状为 (batch_size, channels, depth, height, width)，以生成输出张量。

ceil_mode 用于在计算输出形状时选择向上取整或向下取整。count_include_pad 指示在计算中是否包括或排除填充的输入值。此算子接受数据布局规范。

参数:

data (relax.Expr) – 运算符的输入数据。
pool_size (Union[int, Tuple[int, int]]) – 池化窗口的大小。长度必须为 1 或 3。
strides (Union[int, Tuple[int, int]]) – 池化的步幅。长度必须为 1 或 3。
padding (Union[int, Tuple[int, ...]]) – 池化的填充。长度必须为 1、3 或 6。
dilation (Union[int, Tuple[int, int]]) – 池化的扩张。长度必须为 1 或 3。
ceil_mode (bool) – 一个布尔值，指示是使用向上取整还是向下取整来计算输出形状。通过使用向上取整，输入张量中的每个元素都将被滑动窗口覆盖。
count_include_pad (bool, optional) – 是否包含填充以计算平均值。
layout (str) – 输入的布局（Layout）。
out_layout (Optional[str]) – 输出的布局（Layout）。如果未指定，则与 data_layout 相同

返回:

result – 计算结果。

返回类型:

Expr

tvm.relax.op.nn.nll_loss(predictions: RelaxExpr, targets: RelaxExpr, weights: RelaxExpr | None = None, reduction: str = 'mean', ignore_index: int = -100) → RelaxExpr

负对数似然损失。

output[n, i_1, i_2, …, i_k] = -p * w, 其中 - p = predictions[n, t, i_1, i_2, i_k], - t = targets[n, i_1, i_2, …, i_k], - w = weights[t] if t != ignore_index else 0

result = reduction(output)

参数:

predictions (relax.Expr) – 预测值。应为 (k+2)-D 张量，形状为 (N, C, d_1, d_2, …, d_k)，其中 C 是目标类别的数量。
targets (relax.Expr) – 每个预测的目标值。应为 (k+1)-D 张量，形状为 (N, d_1, d_2, …, d_k)。必须为 int 数据类型。
weights (Optional[relax.Expr]) – 每个目标值的权重。应为 1-D 张量，形状为 (C,)。如果未指定，则视为所有权重均为 1。
reduction (str) – 应用于输出的缩减方法。可能的值为 “mean”、“sum” 和 “none”。
ignore_index (int) – 要忽略的目标值。

返回:

result – 计算结果。

返回类型:

relax.Expr

tvm.relax.op.nn.pad(data: RelaxExpr, pad_width: Tuple[Tuple[int, int], ...], pad_mode: str | None = 'constant', pad_value: float | RelaxExpr | None = 0.0)

填充。

此运算符接受一个张量，并使用指定的值按指定的宽度填充每个轴。

参数:

data (relax.Expr) – 运算符的输入数据
pad_width (Tuple[Tuple[int, int], ...], required) – 填充到每个轴边缘的值的数量，格式为 ((before_1, after_1), …, (before_N, after_N))
pad_mode (Optional[str]) – ‘constant’、‘edge’ 或 ‘reflect’。‘constant’ 使用 constant_value pad_value 填充。‘edge’ 使用输入数组的边缘值填充。‘reflect’ 通过相对于边缘反射值来填充。默认为 ‘constant’。
pad_value (Optional[Union[float, Expr]]) – 用于填充的值。默认为 0。

返回:

result – 计算结果。

返回类型:

relax.Expr

tvm.relax.op.nn.relu(data: RelaxExpr) → RelaxExpr

线性整流单元。

\[\text{ReLU}(x) = \max(x, 0)\]

参数:: data (relax.Expr) – 输入数据
返回:: result – 计算结果。
返回类型:: relax.Expr

tvm.relax.op.nn.rms_norm(data: RelaxExpr, weight: RelaxExpr, axes: int | List[int] = -1, epsilon: float = 1e-05) → RelaxExpr

均方根归一化（Biao Zhang 等人，2019）。将均方根归一化应用于 n 维输入数组。此运算符接受 n 维输入数组，并使用给定的轴对输入进行归一化。

\[out = \frac{data}{\sqrt{mean(data, axis)+\epsilon}} * weight + bias\]

参数:

data (relax.Expr) – 将应用 rms_norm 的输入。
weight (relax.Expr) – 比例因子。
bias (relax.Expr) – 偏移因子。
axes (Union[int, List[int]]) – 应用归一化的轴。
epsilon (float) – 添加到平方均值的小浮点数，以避免被零除。

返回:

result – 计算结果。

返回类型:

relax.Expr

tvm.relax.op.nn.silu(data: RelaxExpr) → RelaxExpr

Sigmoid 线性单元函数

\[\text{SiLU}(x) = x * \text{sigmoid}(x)\]

参数:: data (relax.Expr) – 输入数据
返回:: result – 计算结果。
返回类型:: relax.Expr

Note

输入张量需要具有浮点 dtype

tvm.relax.op.nn.softmax(data: RelaxExpr, axis: int = -1) → RelaxExpr

计算 softmax。

\[\text{softmax}(x)_i = \frac{\exp(x_i)}{\sum_j \exp(x_j)}\]

参数:

data (relax.Expr) – 运算符的输入数据。
axis (int) – 计算 softmax 时要进行求和的轴。如果未指定，则默认为输入张量的最后一个轴。支持负索引。

返回:

result – 计算结果。

返回类型:

relax.Expr

Note

输入张量需要具有浮点 dtype

tvm.relax.op.builtin

Relax 内置运算符。

tvm.relax.op.builtin.alloc_tensor(shape: RelaxExpr, dtype: str | RelaxExpr, runtime_device_index: int | RelaxExpr, storage_scope: str | RelaxExpr = 'global') → Call

构造一个 relax.Call 以分配具有特定形状、dtype、runtime_device_index 的张量。

参数:

shape (Expr) – 要分配的张量的形状。
dtype (Union[str, Expr]) – 要分配的张量的数据类型。
runtime_device_index (Union[int, Expr]) – 设备索引，指示在运行时要在哪个设备上分配张量。索引 -1 保留给主机设备。
storage_scope (Union[str, Expr]) – 要将存储分配到的存储作用域。

返回:

result – 一个 relax relax.Call，它获取已分配的张量。

返回类型:

relax.Call

tvm.relax.op.builtin.stop_lift_params(x: RelaxExpr) → RelaxExpr

一个指示符，表明输入张量的使用者不应被提升到 transform_params 函数

参数:: x (relax.Expr) – 输入数据
返回:: result – 与输入张量相同的输出张量
返回类型:: relax.Expr

tvm.relax.op.ccl

CCL 相关运算符。

tvm.relax.op.ccl.allgather(x, num_workers: int, in_group: bool = True)

AllGather 运算符

参数:

x (relax.Expr) – 输入张量。
num_worker (int) – 从中收集数据的工作节点的数量。
in_group (bool) – 收集操作是全局执行还是默认在组内执行。

返回:

result – allgather 的结果。

返回类型:

relax.Expr

tvm.relax.op.ccl.allreduce(x, op_type: str = 'sum', in_group: bool = True)

Allreduce 运算符

参数:

x (relax.Expr) – 输入张量。
op_type (str) – 要应用于输入数据的缩减操作的类型。现在支持 “sum”、“prod”、“min”、“max” 和 “avg”。
in_group (bool) – 缩减操作是全局执行还是默认在组内执行。

返回:

result – allreduce 的结果。

返回类型:

relax.Expr

tvm.relax.op.ccl.broadcast_from_worker0(x: RelaxExpr) → RelaxExpr

将数据从 worker-0 广播到所有其他工作节点。

参数:: x (relax.Expr) – 要广播的张量。
返回:: result – 相同的张量，已广播到所有其他工作节点。
返回类型:: relax.Expr

tvm.relax.op.ccl.scatter_from_worker0(x: RelaxExpr, num_workers: int, axis: int = 0) → RelaxExpr

从 worker-0 执行 scatter 操作，将给定的缓冲区分块为相等的部分。

参数:

x (relax.Expr) – 要分成相等部分并相应地发送到每个工作节点的缓冲区。
num_worker (int) – 工作节点的数量，即应将给定缓冲区分块成的部分数。
axis (int) – 要分散的张量的维度。默认为 0。

返回:

result – 不同工作节点接收到的分块张量。

返回类型:

relax.Expr

tvm.relax.op.distributed

用于分布式 Relax 的运算符。

tvm.relax.op.distributed.annotate_sharding(input: RelaxExpr, device_mesh: DeviceMesh, placement: Placement) → RelaxExpr

为张量注释分片计划

参数:

input (relax.Expr) – 输入张量。
device_mesh (DeviceMesh) – 分片计划的设备网格
placement (Placement) – 分片计划的放置

返回:

result – 未修改的张量。

返回类型:

relax.Expr

tvm.relax.op.distributed.redistribute(input: RelaxExpr, device_mesh: DeviceMesh, placement: Placement) → RelaxExpr

重分布张量

参数:

input (relax.Expr) – 输入张量。
device_mesh (DeviceMesh) – 重分布后的设备网格
placement (Placement) – 重分布后的放置

返回:

result – 重分布后的张量。

返回类型:

relax.Expr

tvm.relax.op.distributed.call_tir_local_view(gvar: GlobalVar, args: RelaxExpr, out_sinfo: DTensorStructInfo | List[DTensorStructInfo], tir_vars: ShapeExpr | Tuple[PrimExpr] | List[PrimExpr] | None = None) → Call

relax.Call 一个 tir.prim_func 并返回输出。prim_func 应该是一个 worker-local 函数，它实际上在每个 worker 上执行，而不是未分区的函数。此运算符的输出是 DTensor 或 DTensor 元组。

参数:

gvar (GlobalVar) – 指向 tir PrimFunc 的 GlobalVar。
args (Expr) – 输入参数。
out_sinfo (Union[DTensorStructInfo, List[DTensorStructInfo]]) – call_tir 输出的结构信息。它应该是一个或一个 DTensorStructInfo 列表。每一个都表示返回张量的结构信息。
tir_vars (Optional[Union[ShapeExpr, Tuple[PrimExpr], List[PrimExpr]]]) – ShapeExpr，表示调用 func 时要解包的整数元组。如果不使用，则为空

返回:

ret – call_tir_local_view 运算符的调用节点。

返回类型:

relax.Call

tvm.relax.op.distributed.redistribute_replica_to_shard(input: RelaxExpr, num_workers: int, axis: int) → RelaxExpr

沿一个轴将张量切片成几个部分，: 每个工作节点取一部分。需要 input.struct_info.shape[axis] % num_workers == 0。每个工作节点必须具有输入的相同副本。这是 redistribute 运算符的专用版本。

参数:

input (relax.Expr) – 要切片成相等部分的缓冲区。
num_worker (int) – 工作节点的数量，即应将给定缓冲区切片成的部分数。
axis (int) – 要切片的张量的轴。

返回:

result – 每个设备保留的切片张量。

返回类型:

relax.Expr

tvm.relax.op.grad

用于查找 relax 运算符梯度的运算符。

tvm.relax.op.grad.Expr: 别名：RelaxExpr

tvm.relax.op.grad.avg_pool2d_backward(output_grad: RelaxExpr, data: RelaxExpr, pool_size: Tuple[int, int] = (1, 1), strides: Tuple[int, int] = (1, 1), padding: Tuple[int, int, int, int] = (0, 0, 0, 0), dilation: Tuple[int, int] = (1, 1), ceil_mode: bool = False, count_include_pad: bool = False, layout: str = 'NCHW', out_layout: str | None = None) → RelaxExpr

relax.nn.avg_pool2d 的反向运算符。除 output_grad 外的所有参数均与 relax.nn.avg_pool2d 相同。返回关于 data 的梯度。

参数:: output_grad (relax.Expr) – 关于 avg_pool2d 结果的梯度。
返回:: result – 关于 data 的梯度。
返回类型:: relax.Expr

tvm.relax.op.grad.end_checkpoint(input: RelaxExpr) → RelaxExpr

标记检查点阶段的结束。请参阅 tvm.relax.op.grad.start_checkpoint。

参数:: input (relax.Expr) – 检查点阶段的输出。
返回:: result – 与输入相同的张量。
返回类型:: relax.Expr

tvm.relax.op.grad.max_pool2d_backward(output_grad: RelaxExpr, data: RelaxExpr, pool_size: Tuple[int, int] = (1, 1), strides: Tuple[int, int] = (1, 1), padding: Tuple[int, int, int, int] = (0, 0, 0, 0), dilation: Tuple[int, int] = (1, 1), ceil_mode: bool = False, count_include_pad: bool = False, layout: str = 'NCHW', out_layout: str | None = None) → RelaxExpr

relax.nn.max_pool2d 的反向运算符。除 output_grad 外的所有参数均与 relax.nn.max_pool2d 相同。返回关于 data 的梯度。

参数:: output_grad (relax.Expr) – 关于 max_pool2d 结果的梯度。
返回:: result – 关于 data 的梯度。
返回类型:: relax.Expr

tvm.relax.op.grad.nll_loss_backward(output_grad: RelaxExpr, predictions: RelaxExpr, targets: RelaxExpr, weights: RelaxExpr | None = None, reduction: str = 'mean', ignore_index: int = -100) → RelaxExpr

relax.nn.nll_loss 的反向运算符。除 output_grad 外的所有参数均与 relax.nn.nll_loss 相同。返回关于 predictions 的梯度。

参数:: output_grad (relax.Expr) – 关于 nll_loss 结果的梯度。
返回:: result – 关于 predictions 的梯度。
返回类型:: relax.Expr

tvm.relax.op.grad.no_grad(input: RelaxExpr) → RelaxExpr

关于输入的无梯度虚拟运算符。

参数:: input (relax.Expr) – 对应的输入张量。
返回:: result – 关于输入的无梯度表示。
返回类型:: relax.Expr

tvm.relax.op.grad.start_checkpoint(input: RelaxExpr) → RelaxExpr

标记检查点阶段的开始。start_checkpoint 和 end_checkpoint 之间的计算将被标记为检查点阶段。

检查点阶段不保存所有中间激活，而是在反向传播过程中重新计算它们，而不是存储整个计算图的所有中间激活以用于计算反向传播。

例如，` a = relax.Var("a", relax.TensorStructInfo((2, 2), "float32")) b = relax.Var("b", relax.TensorStructInfo((2, 2), "float32")) c = a * 2 d = b * 2 c_cp = start_checkpoint(c) d_cp = start_checkpoint(d) e = c_cp + d_cp e_out = end_checkpoint(e) ` 然后将在反向传播阶段重新计算 e。

有关更多信息，请参阅 tvm.relax.transform.Gradient、tvm.relax.testing.nn.checkpoint、tvm.relax.op.grad.end_checkpoint。

参数:: input (relax.Expr) – 标记检查点阶段输入的张量。
返回:: result – 与输入相同的张量。
返回类型:: relax.Expr

tvm.relax.op.grad.take_backward(output_grad: RelaxExpr, x: RelaxExpr, indices: RelaxExpr, axis: int | None = None) → RelaxExpr

relax.take 的反向运算符。除 output_grad 外的所有参数均与 relax.take 相同。返回关于 x 的梯度。

参数:: output_grad (relax.Expr) – 关于 take 结果的梯度。
返回:: result – 关于 x 的梯度。
返回类型:: relax.Expr

tvm.relax.op.image

图像运算符。

tvm.relax.op.image.resize2d(data: RelaxExpr, size: RelaxExpr | int | PrimExpr | Tuple[int | PrimExpr], roi: float | Tuple[float] | None = None, layout: str = 'NCHW', method: str = 'linear', coordinate_transformation_mode: str = 'half_pixel', rounding_method: str = 'round', cubic_alpha: float = -0.5, cubic_exclude: int = 0, extrapolation_value: float = 0.0, out_dtype: str | DataType | None = None) → RelaxExpr

图像 resize2d 运算符。

此运算符将数据作为输入，并对给定的缩放因子进行 2D 缩放。在默认情况下，当 data_layout 为 NCHW 且数据形状为 (n, c, h, w) 时，输出的形状将为 (n, c, size[0], size[1])

method 指示计算输出值时要使用的算法，method 可以是 (“linear”, “nearest_neighbor”, “cubic”) 之一

参数:

data (relax.Expr) – 运算符的输入数据。
size (Union[Expr, PrimExprLike, Tuple[PrimExprLike]]) – 图像将被调整为的输出大小。如果指定为列表，则其长度必须为 1 或 2。如果指定为 Expr，则其 ndim 必须为 2。
roi (Optional[Union[float, Tuple[float]]]) – 用于裁剪输入图像的感兴趣区域。预期大小为 4，格式为 [start_h, start_w, end_h, end_w]。仅当 coordinate_transformation_mode 为 tf_crop_and_resize 时使用。
layout (str) – 输入的布局（Layout）。
method (str) – 要使用的缩放方法 [nearest_neighbor, linear, cubic]。
coordinate_transformation_mode (str) – 描述如何将调整大小的张量中的坐标转换为原始张量中的坐标。定义可以在 topi/image/resize.py 中找到。[half_pixel, align_corners, asymmetric, pytorch_half_pixel, tf_half_pixel_for_nn, 和 tf_crop_and_resize]。
rounding_method (str) – 指示如何在 nearest_neighbor 方法中查找“最近”像素 [round, floor, ceil]
cubic_alpha (float) – 双三次插值的样条系数
cubic_exclude (int) – 在双三次插值期间排除图像外部的标志
extrapolation_value (float) – 当 roi 在图像外部时使用的填充值
out_dtype (Optional[Union[str, DataType]]) – 输出张量的数据类型。如果未指定，则如果未指定，则输出将具有与输入相同的数据类型。

返回:

result – 调整大小后的结果。

返回类型:

relax.Expr

tvm.relax.op.memory

Relax 内存原语。

tvm.relax.op.memory.alloc_storage(size: RelaxExpr, virtual_device_index: int | RelaxExpr, storage_scope: str | RelaxExpr, dtype: str | RelaxExpr) → Call

构造一个 relax.Call，以分配具有特定大小、virtual_device_index、storage_scope 和 dtype 的存储。

参数:

size (Expr) – 要分配的存储的大小。
virtual_device_index (Union[int, Expr]) – 虚拟设备索引，指示要在哪个设备上分配存储。索引 -1 保留给主机设备。
storage_scope (Union[str, Expr]) – 要将存储分配到的存储作用域。
dtype (Union[str, Expr]) – 要分配的存储的数据类型。

返回:

result – 一个 relax relax.Call，它获取分配的存储。

返回类型:

relax.Call

tvm.relax.op.memory.alloc_tensor(storage: RelaxExpr, offset: int | RelaxExpr, shape: RelaxExpr, dtype: str | RelaxExpr) → Call

构造一个 relax.Call，以在从给定偏移量开始的特定存储上分配张量。

参数:

storage (Expr) – 要在其中分配张量的存储。
offset (Union[int, Expr]) – 用于分配张量的存储偏移量。
shape (Expr) – 要分配的张量的形状。
dtype (Union[str, Expr]) – 要分配的张量的数据类型。

返回:

result – 一个 relax relax.Call，它获取已分配的张量。

返回类型:

relax.Call

tvm.relax.op.memory.kill_storage(storage: RelaxExpr) → Call

构造一个 relax.Call 以销毁存储。

参数:: storage (Expr) – 要销毁的存储。
返回:: result – 一个 relax relax.Call，用于销毁存储。
返回类型:: relax.Call

tvm.relax.op.memory.kill_tensor(tensor: RelaxExpr) → Call

构造一个 relax.Call 以销毁张量。

参数:: tensor (Expr) – 要销毁的张量。
返回:: result – 一个 relax relax.Call，用于销毁张量。
返回类型:: relax.Call

提供现有张量的视图

视图可能具有不同的形状、不同的数据类型，并且可能相对于源数组以偏移量开始。

无论使用这些选项的哪种组合，视图都绝不能访问输入 data 数组无法访问的内存。即使 data 数组本身是共享后备数组的视图，此限制也适用。

参数:

data (relax.Expr) – 运算符的输入数据。
shape (Optional[Union[Sequence[PrimExprLike], Expr]]) – 目标形状。应为 relax.ShapeExpr，或可转换为 relax.ShapeExpr 的集合。
dtype (Optional[Expr]) – 目标数据类型。应为 relax.ShapeExpr，或可转换为 relax.ShapeExpr 的集合。
relative_byte_offset (Optional[Expr]) – 输出 NDArray 的偏移量，相对于 data 的字节偏移量。如果为 None，则视图的偏移量与 data 的偏移量相同。

返回:

result – 张量视图

返回类型:

relax.Expr

tvm.relax.op.memory.ensure_zero_offset(data: RelaxExpr) → RelaxExpr

确保张量的 elem_offset == 0。如有必要，将进行复制。

参数:

data (relax.Expr) – 输入张量
结果
-------
result (relax.Expr) – elem_offset == 0 的张量

tvm.relax.op.op_attrs

用于 Relax 运算符的属性节点

class tvm.relax.op.op_attrs.CallTIRWithGradAttrs: call_tir_with_grad 运算符中使用的属性

class tvm.relax.op.op_attrs.InitAttrs: full/full_like、ones/ones_like 和 zeros/zeros_like 运算符中使用的属性

class tvm.relax.op.op_attrs.TriluAttrs: tril 和 triu 运算符中使用的属性

class tvm.relax.op.op_attrs.AstypeAttrs: astype 运算符中使用的属性

class tvm.relax.op.op_attrs.TakeAttrs: take 运算符中使用的属性

class tvm.relax.op.op_attrs.StridedSliceAttrs: strided_slice 运算符中使用的属性

class tvm.relax.op.op_attrs.MatmulAttrs: matmul 运算符的属性

class tvm.relax.op.op_attrs.Conv2DAttrs: nn.conv2d 的属性

class tvm.relax.op.op_attrs.Conv3DAttrs: nn.conv3d 的属性

class tvm.relax.op.op_attrs.Conv2DTransposeAttrs: nn.conv2d_transpose 的属性

class tvm.relax.op.op_attrs.Pool2DAttrs: nn.max_pool2d 的属性

class tvm.relax.op.op_attrs.AdaptivePool2DAttrs: 2d 自适应池化运算符的属性

class tvm.relax.op.op_attrs.SoftmaxAttrs: nn.softmax 的属性

class tvm.relax.op.op_attrs.BatchNormAttrs: batch_norm 运算符中使用的属性

class tvm.relax.op.op_attrs.LayerNormAttrs: layer_norm 运算符中使用的属性

class tvm.relax.op.op_attrs.DropoutAttrs: dropout 运算符的属性

class tvm.relax.op.op_attrs.StatisticalAttrs: statistical 运算符中使用的属性

class tvm.relax.op.op_attrs.ConcatAttrs: concat 运算符的属性

class tvm.relax.op.op_attrs.ExpandDimsAttrs: expand_dims 运算符的属性

class tvm.relax.op.op_attrs.PermuteDimsAttrs: permute_dims 运算符的属性

class tvm.relax.op.op_attrs.SortAttrs: sort 运算符的属性

class tvm.relax.op.op_attrs.ArgsortAttrs: argsort 运算符的属性

class tvm.relax.op.op_attrs.SplitAttrs: split 运算符中使用的属性

class tvm.relax.op.op_attrs.SqueezeAttrs: squeeze 运算符的属性

class tvm.relax.op.op_attrs.LayoutTransformAttrs: layout_transform 运算符中使用的属性

class tvm.relax.op.op_attrs.Resize2DAttrs: image resize2d 运算符中使用的属性

class tvm.relax.op.op_attrs.ArgmaxArgminAttrs: argmax/argmin 运算符的属性

class tvm.relax.op.op_attrs.RepeatAttrs: repeat 运算符的属性

class tvm.relax.op.op_attrs.TileAttrs: tile 运算符的属性

class tvm.relax.op.op_attrs.ScanopAttrs: scan 运算符的属性

class tvm.relax.op.op_attrs.TopKAttrs: topk 运算符的属性

class tvm.relax.op.op_attrs.EinsumAttrs: einsum 运算符的属性

class tvm.relax.op.op_attrs.FlipAttrs: flip 运算符的属性