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 操作的调用(在 Python 中,assert 是保留字,因此名称必须不同)。
- 参数:
- 返回:result:一个调用 Relax 断言操作的 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 (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 一个目标传递风格的打包函数并返回输出。
注意:被调用的函数被认为是纯函数(除了修改指定的输出参数)。如果函数确实产生了其他副作用,那么编译器可能会最终移除、重新排序或重复这些效果——不提供任何保证。
- 参数:
- func (Union[str,Expr]):目标传递风格的函数,可以是 ExternFunc。
- args (Expr):输入参数。
- out_sinfo (Union[TensorStructInfo,List[TensorStructInfo]]*):call_dps_packed 输出的结构信息。它应该是一个 TensorStructInfo 或 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
构造一个调用打包函数的调用,该函数会消耗其部分参数「就地」并返回被修改的参数(别名),但应被视为其他方面是纯函数。inplace_indices 参数指示哪些输出是被修改的参数。
生成的调用将具有与直接调用打包函数相同的语义。
注意:这应该用于用户知道使用这些参数调用打包函数实际上不会引起任何其他副作用的情况。如果它用于一个确实会导致其他副作用的调用,那么编译器可能会最终删除、重新排序或重复该调用,并且不保证调用方产生的任何副作用。
警告:此运算符在类型系统中被视为纯操作,即使它执行了副作用(修改某些参数)。因此,用户必须确保它被安全使用(即,修改后的参数在修改后不应处于活动状态,它们不应在修改后别名化值)。
- 参数:
- 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
构建一个调用打包函数的调用,即使打包调用通常不被视为纯函数。
生成的调用将具有与直接调用打包函数相同的语义。
注意:这应该用于用户知道使用这些参数调用打包函数实际上不会产生任何副作用的情况。如果用于一个确实会导致副作用的调用,那么编译器可能会最终移除、重新排序或重复该调用,并且对被调用方的任何副作用不做任何保证。
- 参数:
- 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]]):调用_tir 的输出结构信息。它应该是一个 TensorStructInfo 或 TensorStructInfo 的列表。每个元素表示一个返回张量的结构信息。
- tir_vars (Optional[Union[ShapeExpr, Tuple[PrimExpr],List[PrimExpr]]]):表示调用 func 时解包的整数元组的 ShapeExpr。如果未使用则为 null。
- 返回:ret:call_tir 运算符的调用节点。
- 返回类型:relax.Call。
tvm.relax.op.call_tir_inplace(gvar:GlobalVar, args:RelaxExpr, inplace_indices:int|List[int], out_sinfo:TensorStructInfo|List[TensorStructInfo], tir_vars:ShapeExpr|Tuple[PrimExpr] |List[PrimExpr] |None= None) → Call
relax 调用 TIR PrimFunc 并返回结果,在指定位置执行计算(基于 inplace_indices 参数;输出将与通过 in-place 索引选择的输入别名)。
警告:此运算符在类型系统中被视为纯操作,但实际上会修改由 inplace_indices 指定的参数。此运算符不应直接使用,而应由检查是否安全进行原地操作(即,作为输出的参数没有被别名化或调用 call_tir_inplace 后仍然活跃)的 passes 插入。
仅应出于测试目的对此运算符进行直接调用。
- 参数:
- gvar ( GlobalVar ):GlobalVar 引用 TIR PrimFunc。
- 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 时需要解包的整数元组。若未使用则为 null。
- 返回: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 调用 tir.prim_func 并返回输出。这个内建函数会将 te 梯度函数(通过 te_grad_name 指向)绑定到 call_tir_with_grad 节点。梯度传递过程会调用这个 te 梯度函数。
- 参数:
- 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**]]]):表示调用 func 时要解包的整数元组的 ShapeExpr。如果未使用则为 null
- 返回: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 Call,在运行时将打印值。
- 返回类型: Expr。
tvm.relax.op.register_gradient(op_name:str, fgradient:Callable[[Var,Call,Var,BlockBuilder],List[RelaxExpr]] = None, level:int= 10)
为 relax 运算符注册运算符梯度函数。
- 参数:
tvm.relax.op.shape_of(expr:RelaxExpr) → RelaxExpr
获取张量的形状。
- 参数:
- expr (Expr):输入的 Expr。
- 返回:result:一个 relax.Call,它获取输入的形状。
- 返回类型: Expr。
tvm.relax.op.shape_to_tensor(expr:RelaxExpr) → RelaxExpr
将形状转换为张量 expr。 :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
按位 AND :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
按位 OR :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.log_add_exp(x1:RelaxExpr, x2:RelaxExpr) → RelaxExpr
计算输入的指数和的对数,逐元素进行。
- 参数:
- x1 (Expr):第一个输入张量。
- x2 (Expr):第二个输入张量。
- 返回: The element-wise log-sum-exp of x1 and x2。
- 返回类型: 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。
tvm.relax.op.arange(start:int|PrimExpr|PrimValue, end:int|PrimExpr|PrimValue|None= None, step:int|PrimExpr|PrimValue= 1, dtype:str| dtype |None= None) → RelaxExpr
构造一个元素均匀分布的张量。
- 参数:
- 返回:result:结果张量。
- 返回类型: relax.Expr。
tvm.relax.op.full(shape:Tuple[int|PrimExpr] |RelaxExpr, fill_value:RelaxExpr, dtype:str| dtype |None= None) → RelaxExpr
用标量值填充数组。
- 参数:
- 返回:result:结果张量。
- 返回类型: relax.Expr。
tvm.relax.op.full_like(x:RelaxExpr, fill_value:RelaxExpr, dtype:str| dtype |None= None) → RelaxExpr
构造一个张量,其特征为:- 形状与输入数据张量的形状相同,- 其值被输入标量 fill_value 填充。
- 参数:
- x (relax.Expr):输入张量,提供形状,当 dtype 字段未指定时提供 dtype。
- fill_value (relax.Expr):填充的值。必须是一个标量张量。
- dtype (Optional*[Union[str, DataType]]):创建的张量的数据类型。如果未给出 dtype,默认将使用输入张量的 dtype。
- 返回:result:结果张量。
- 返回类型: relax.Expr。
tvm.relax.op.hamming_window(window_size, periodic, alpha, beta, dtype)
Hamming 窗函数。
- 参数:
- 返回:ret:结果张量。
- 返回类型: relax.Expr。
tvm.relax.op.ones(shape:Tuple[int|PrimExpr] |RelaxExpr, dtype:str| dtype) → RelaxExpr
构建一个所有元素都为 1 的张量,具有输入的形状和数据类型。
- 参数:
- 返回:result:结果张量。
- 返回类型: relax.Expr。
tvm.relax.op.ones_like(x:RelaxExpr, dtype:str| dtype |None= None) → RelaxExpr
构造一个全为 1 的张量,其形状与输入张量的形状相同。
- 参数:
- 返回:result:结果张量。
- 返回类型: relax.Expr。
tvm.relax.op.eye(n:int|PrimExpr|PrimValue, m:int|PrimExpr|PrimValue|None= None, k:int|PrimExpr|PrimValue= 0, dtype:str| dtype = 'float32') → RelaxExpr
构造一个对角线上为 1,其余位置为 0 的 2-D 张量。
- 参数:
- 返回:result:结果张量。
- 返回类型: relax.Expr。
tvm.relax.op.eye_like(x:RelaxExpr, k:int|PrimExpr|PrimValue= 0, dtype:str| dtype |None= None) → RelaxExpr
返回一个二维张量,其对角线元素为 1,其余元素为 0,形状与输入张量相同。
- 参数:
- 返回: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) :指示要置零元素的 diagonals 以下索引。如果 k = 0,对角线是主对角线。如果 k < 0,对角线在主对角线下方。如果 k > 0,对角线在主对角线上方。
- 返回:ret :结果张量。
- 返回类型: relax.Expr。
tvm.relax.op.zeros(shape:Tuple[int|PrimExpr] |RelaxExpr, dtype:str| dtype) → RelaxExpr
构建一个全为零的张量,其输入形状和 dtype。
- 参数:
- 返回:result :结果张量。
- 返回类型: relax.Expr
tvm.relax.op.zeros_like(x:RelaxExpr, dtype:str| dtype |None= None) → RelaxExpr
构造一个全为零的张量,其形状与输入张量的形状相同。
- 参数:
- x (relax.Expr) :输入张量,当 dtype 字段未指定时,提供形状和 dtype。
- dtype (Optional[Union[str,DataType]]) : 创建的张量的数据类型。如果未给出 dtype,默认将使用输入张量的数据类型。
- 返回:result :结果张量。
- 返回类型: relax.Expr。
tvm.relax.op.astype(x:RelaxExpr, dtype:str| dtype) → RelaxExpr
将输入张量转换为指定的数据类型。
- 参数:
- x (relax.Expr) :运算符的输入数据。
- dtype (Union[str,DataType]) :目标数据类型。
- 返回:result :转换结果。
- 返回类型: relax.Expr。
tvm.relax.op.wrap_param(data:RelaxExpr, dtype:str| dtype = 'float32') → RelaxExpr
如果输入数据的类型与给定的类型不同,则将模型参数输入张量转换为数据类型。 :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) **:**指定步长值,可以是负数,此时输入张量将在该特定轴上反转。如果未指定,默认为与轴数相同长度的全为 1 的列表。
- 返回:ret :切片结果。
- 返回类型: relax.Expr。
dyn_strided_slice 要求输入的 begin、end 和 strides 的长度与数据张量的 rank 相同。
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]****]) :指定步长值,可以是负数,此时输入张量将在该特定轴上反转。如果未指定,默认为与轴数相同长度的全为 1 的列表。
- assume_inbound (bool) :是否假设索引在范围内。如果设置为 false,超出范围的索引将被裁剪到边界。
- 返回:ret :切片结果。
- 返回类型: relax.Expr。
strided_slice 要求输入的 begin、end 和 strides 的长度与 axes 相同。
:::
tvm.relax.op.take(x:RelaxExpr, indices:RelaxExpr, axis:int|None= None) → RelaxExpr
从一个张量沿轴取出元素。它的语义与 numpy.take(https://numpy.org/doc/stable/reference/generated/numpy.take.html)基本相同,可以涵盖 torch.take(https://pytorch.org/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)
对数据进行爱因斯坦求和约定求值。
- 参数:
- 返回:result :einsum op 的输出。
- 返回类型: relax.Expr。
tvm.relax.op.linear(data:RelaxExpr, weight:RelaxExpr, bias:RelaxExpr|None= None, out_dtype:str| dtype |None= None) → RelaxExpr
对输入数据应用线性变换:y = xA^T + b。
- 参数:
- data (relax.Expr) : 输入数据。
- weight (relax.Expr) :权重张量。
- bias (Optional[Expr]) :偏置张量。
- out_dtype (Optional*[Union[str,DataType]]) :乘法结果的数值类型。如果未指定,输出类型将与输入类型相同。
注意
Relax 不将线性运算符视为原始运算符,而是通过组合转置、矩阵乘法和加法运算符来实现它。
- 返回:result : 计算结果。
- 返回类型: relax.Expr。
tvm.relax.op.matmul(x1:RelaxExpr, x2:RelaxExpr, out_dtype:str| dtype |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 将与输入 dtype 相同。
- 返回:result : 计算结果。
- 返回类型: relax.Expr
tvm.relax.op.outer(x1:RelaxExpr, x2:RelaxExpr) → RelaxExpr
计算两个输入表达式的外积。
- 参数:
- x1 (relax.Expr) :第一个输入表达式。
- x2 (relax.Expr) :第二个输入表达式。
注意
该操作计算两个表达式之间的外积,结果是一个张量,其中每个元素是来自 x1 和 x2 的元素的乘积。它在张量和矩阵运算中常用,用于将低维输入扩展为高维表示。
- 返回:result :表示外积的结果表达式。
- 返回类型: relax.Expr。
tvm.relax.op.broadcast_to(x:RelaxExpr, shape:Tuple[int|PrimExpr] |RelaxExpr) → RelaxExpr
将张量广播到指定形状。
- 参数:
- 返回:result :广播后的张量。
- 返回类型: relax.Expr。
tvm.relax.op.collapse_sum_like(data:RelaxExpr, collapse_target:RelaxExpr) → RelaxExpr
返回数据到 collapse_target 的形状的求和。
详情请参阅 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
返回数据折叠到给定形状的求和结果。
collapse_sum_to 是 tvm.relax.op.broadcast_to 和其他广播运算符在自动微分过程中的反向运算符。
我们预期 data 是某个广播操作中广播某个给定形状张量的结果。因此,给定形状和 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
沿给定轴连接输入张量。
- 参数:
- 返回:result : 连接的张量。
- 返回类型: relax.Expr。
tvm.relax.op.expand_dims(x:RelaxExpr, axis:int|List[int]) → RelaxExpr
在 axis 指定的位置插入新的轴。
- 参数:
- 返回: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
使用 updates 中的值更新由 indices 定义的位置处的数据。
- 参数:
- 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]
tvm.relax.op.index_put(data:RelaxExpr, indices:RelaxExpr|Tuple[RelaxExpr], values:RelaxExpr, accumulate:bool= False) → RelaxExpr
此操作使用来自 values 的对应值更新 data 中由 indices 指定的位置。indices 是一个张量元组,其中每个张量对应 data 中的一个维度。当 accumulate 为 True 时,操作执行累积(加法)而不是替换。reduction 参数允许指定不同的缩减操作。 :param data: 要修改的输入张量 :type data: relax.Expr :param indices: 指定更新位置的索引张量元组(每个维度一个) :type indices: Union[Expr, Tuple[Expr]] :param values: 要放置在指定索引处的值 :type values: relax.Expr :param accumulate: 是否累积(加)值而不是替换(默认:False) :type accumulate: bool
- 返回:result :一个与 data 形状相同但指定位置已更新的新张量。
- 返回类型: relax.Expr。
示例
# inputs
data = torch.zeros(3, 3)
indices = (torch.tensor([0, 2]), torch.tensor([1, 1]))
values = torch.tensor([1.0, 2.0])
# output
output = [
[0.0, 1.0, 0.0],
[0.0, 0.0, 0.0],
[0.0, 2.0, 0.0],
]
# with accumulate=True
output = [
[0.0, 1.0, 0.0],
[0.0, 0.0, 0.0],
[0.0, 3.0, 0.0],
]
tvm.relax.op.index_tensor(data:RelaxExpr, indices:RelaxExpr|List[RelaxExpr]) → RelaxExpr
高级张量索引(NumPy/PyTorch 风格)。
给定 k 个索引张量 indices = (I0, I1, …, Ik‑1)
,此运算符从 data
中选择元素,就像在 NumPy/PyTorch 中编写了 data[I0, I1, …, Ik‑1]
一样:
所有索引张量必须具有整数 dtype。
它们的形状按照常规的 NumPy 方式广播到一个公共形状 B
。
结果形状是 B + data.shape[k:]
(即广播形状后跟 data
中未索引的剩余轴)。
在编译时,Relax 检查索引张量的数量 k
是否不超过 data.ndim
,dtype 是否为整数,以及形状是否一致(广播兼容)。
- 参数:
- data (relax.Expr)
:
要索引的输入张量。 - indices (Union*[relax.Expr,* List*[relax.Expr**]]*) :一个包含索引张量的元组表达式,或一个将被自动提升为元组表达式的 Python
list
/tuple
。每个张量必须具有整数 dtype。
- data (relax.Expr)
- 返回:result :高级索引后得到的张量。其 dtype 等于
data.dtype
。 - 返回类型: relax.Expr。
示例
import numpy as np
import tvm.relax as R
x = R.const(np.arange(9).reshape(3, 3).astype("float32"))
row = R.const(np.array([0, 2])) # shape (2,)
col = R.const(np.array([1, 0])) # shape (2,)
y = R.index_tensor(x, [row, col])
# y.shape == (2,) ; y == [1., 6.]
# Broadcasting: row : (2,1), col : (1,3) → B = (2,3)
row = R.const(np.array([[0],[1]]))
col = R.const(np.array([[0,1,2]]))
z = R.index_tensor(x, [row, col])
# z.shape == (2,3)
tvm.relax.op.meshgrid(tensors:RelaxExpr|List[RelaxExpr], indexing:str|None= 'ij') → RelaxExpr
从输入张量生成坐标网格。
- 参数:
- tensors (Union[relax.Expr, List[relax.Expr]]) :一个元组类型的 Expr,包含 1D 张量(或提升为 1D 的标量),用于生成坐标网格,或此类张量的列表。
- indexing (Optional[str]) : 索引模式,为“ij”(矩阵索引)或“xy”(笛卡尔索引)。默认为“ij”。
- 返回:result : 一个表示坐标网格的张量元组。
- 返回类型: relax.Expr。
tvm.relax.op.layout_transform(x:RelaxExpr, index_map:Callable|IndexMap, pad_value:int|float|PrimValue|None= None, axis_separators:int| axis_separator |None= None, input_axis_separators:int| axis_separator |None= None)
修改张量的布局。
- 参数:
- 返回: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) :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
重复数组中的元素。
- 参数:
- 返回: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
通过使用输入维度的余数来推断输出形状的维度,同时保持新数组的大小与输入数组相同。shape 最多可以有一个维度为-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。
-1
推理仅在编译时执行。也就是说,在任何情况下,如果 -1
的维度长度不能在编译时推断,将会抛出错误。
tvm.relax.op.scatter_elements(data:RelaxExpr, indices:RelaxExpr, updates:RelaxExpr, axis:int= 0, reduction:str= 'update')
ONNX 风格的散布元素。此操作根据 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。
- 参数:
- 返回: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.slice_scatter(input_tensor:RelaxExpr, src:RelaxExpr, start, end, step, axis=0)
将 src 张量的值嵌入到 input 的指定维度中。
- 参数:
- 返回:result :与 data 相同形状的计算结果张量。
- 返回类型: relax.Expr。
tvm.relax.op.split(x:RelaxExpr, indices_or_sections:int|List[int|PrimExpr], axis:int= 0) → RelaxExpr
沿轴通过分段或索引分割输入张量。
如果 indices_or_sections 是一个整数,输入将沿给定轴均匀分割(如果可能)。如果沿给定维度的张量大小不能被整数整除,最后一个分段将更小。
如果 indices_or_sections 是一个由整数或 PrimExpr 混合组成的元组,则其中的条目指示沿轴分割数组的索引。
- 参数:
- 返回:ret :计算结果。
- 返回类型: relax.Expr。
tvm.relax.op.squeeze(x:RelaxExpr, axis:int|List[int] |None= None) → RelaxExpr
在数组中压缩轴。
- 参数:
- 返回:result :压缩后的结果。
- 返回类型: relax.Expr。
tvm.relax.op.stack(tensors:RelaxExpr|List[RelaxExpr], axis:int= 0) → RelaxExpr
沿着新轴堆叠输入张量。
- 参数:
- tensors (Union*[relax.Expr,* List*[relax.Expr**]]*) :一个包含待堆叠张量的元组类型的 Expr,或张量列表。所有输入张量必须具有相同的形状。
- axis (int) :结果张量中输入张量将被堆叠的轴。负值会循环。默认为 0。
- 返回:result :与输入张量相比,堆叠的张量多了一个维度。
- 返回类型: relax.Expr。
tvm.relax.op.tile(data:RelaxExpr, repeats:int|Tuple[int] |List[int]) → RelaxExpr
通过 repeats 指定的次数重复 data 来构造一个数组。
如果 repeats 的长度为 l,data 的维度为 d,则结果将具有 max(l, d)的维度。
如果 d < l,数据将通过在前面添加新的轴来提升为 l 维。因此,形状为 (3,) 的张量在 2 维复制时会提升为 (1, 3),在 3 维复制时会提升为 (1, 1, 3)。如果这不是期望的行为,请在调用此函数之前手动将数据提升为 d 维。
如果 d > l,reps 将通过在其前面添加 1 来提升为长度 d。因此,对于形状为 (2, 3, 4, 5) 的数据,一个 reps 为 (2, 2) 将被视为 (1, 1, 2, 2)。
- 参数:
- 返回: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)。
- 参数:
- 返回: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)。
- 参数:
- 返回:result :计算结果。
- 返回类型: tvm.relax.Expr。
tvm.relax.op.multinomial_from_uniform(prob:RelaxExpr, uniform_sample:RelaxExpr, sample_indices:RelaxExpr, dtype:str= 'int64') → RelaxExpr
返回一个张量,其中每一行包含从位于张量 prob 相应行的多项式概率分布中采样的索引。
注意
为了更好的 CPU 性能,请使用‘vm.builtin.multinomial_from_uniform’。为了获得准确的结果,请确保概率在 0 和 1 之间且总和为 1。
- 参数:
- prob (relax.Expr) :D 张量,表示概率分布。每一行代表一个批次在词汇表上的分布,其中:值范围在 [0, 1] 之间,表示每个词汇项的概率。每行的值之和为 1,形成一个有效的分布。
- uniform_sample (relax.Expr) :D 张量。值范围在 0 到 1 之间,表示均匀采样的概率。
- sample_indices (relax.Expr) :D 张量,指示要采样的具体概率分布。sample_indices[i] 的值决定了第 i 个词应该从第 sample_indices[i] 个概率分布中采样。例如,如果有 3 个不同的概率分布,并且要求从每个分布中采样 2、3 和 4 个词,那么 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
计算给定轴上张量元素的 argmax。
- 参数:
- 返回:result : 计算结果。
- 返回类型: relax.Expr。
tvm.relax.op.argmin(x:RelaxExpr, axis:int|None= None, keepdims:bool= False) → RelaxExpr
计算给定轴上张量元素的最小值。
- 参数:
- 返回:result : 计算结果。
- 返回类型: relax.Expr。
tvm.relax.op.where(condition:RelaxExpr, x1:RelaxExpr, x2:RelaxExpr) → RelaxExpr
根据条件值从输入张量中选择元素。
对于给定位置,如果条件为 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 :一个包含非零元素索引的二维张量。
- 返回类型: relax.Expr。
该函数等价于 onnx.nonzero。
示例
x = [[0, 1],
[2, 0]]
nonzero(x) = [[0, 1],
[1, 0]]
tvm.relax.op.unique(x:RelaxExpr, sorted:bool|RelaxExpr= True, return_index:bool|RelaxExpr= False, return_inverse:bool|RelaxExpr= False, return_counts:bool|RelaxExpr= False, axis:int|RelaxExpr|None= None) → RelaxExpr
查找给定张量中的唯一元素。此外,它还可以选择性地返回 - 输入张量中给出唯一值的索引; - 重建输入张量的唯一张量的索引; - 每个唯一值在输入张量中出现的次数。
- 参数:
- x (relax.Expr) :输入张量。
- sorted (Union[bool,Expr]) : 是否在返回输出前对唯一元素进行升序排序。
- return_index (Union[bool,Expr]) :是否返回一个额外的张量,包含唯一张量中元素在原始输入中的索引。
- return_inverse (Union[bool,Expr]) :是否返回一个额外的张量,包含原始输入中的元素在返回的唯一列表中的索引。
- return_counts (Union[bool,Expr]) :是否返回一个额外的张量,其中包含每个唯一元素的数量。
- axis (Optional) :应用于唯一值的维度。如果未指定,则返回扁平化输入的唯一值。
- 返回:ret :创建的 relax 调用,
- 返回类型: relax.Expr
tvm.relax.op.argsort(data:RelaxExpr, axis:int= -1, descending:bool= False, dtype:str= 'int32')
沿给定轴进行排序,并返回一个与输入数组形状相同的索引数组,该数组按排序顺序索引数据。
- 参数:
- 返回:out :与 data 形状相同的张量。
- 返回类型: relax.Expr,
tvm.relax.op.sort(x:RelaxExpr, axis:int= -1, 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')
获取输入张量沿给定轴的 top k 元素。
ret_type 指定返回类型,可以是("both","values","indices")之一。
- 参数:
- 返回:out : 计算结果。
- 返回类型: relax.Expr or List[relax.Expr]。
tvm.relax.op.cumprod(data:RelaxExpr, axis:int|None= None, dtype:str| dtype |None= None, exclusive:bool= False)
Numpy 风格的累积乘积操作。返回沿给定轴的元素累积乘积。
- 参数:
- 返回:result :结果与 data 的大小相同,如果 axis 不为 None,则形状也与 data 相同。如果 axis 为 None,结果是一个 1 维数组。
- 返回类型: 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| dtype |None= None, exclusive:bool= False)
Numpy 风格的累积和操作。返回沿给定轴的元素的累积包含和。
- 参数:
- 返回:result :结果与 data 的大小相同,如果 axis 不为 None,则形状也与 data 相同。如果 axis 为 None,结果是一个 1 维数组。
- 返回类型: 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
计算给定轴上张量元素的最大值。
- 参数:
- 返回:result :计算结果。
- 返回类型: relax.Expr。
tvm.relax.op.mean(x:RelaxExpr, axis:int|List[int] |None= None, keepdims:bool= False) → RelaxExpr
计算给定轴上张量元素的平均值。
- 参数:
- 返回:result :计算结果。
- 返回类型: relax.Expr。
tvm.relax.op.min(x:RelaxExpr, axis:int|List[int] |None= None, keepdims:bool= False) → RelaxExpr
计算在给定轴上张量元素的最小值。
- 参数:
- 返回:result :计算结果。
- 返回类型: relax.Expr。
tvm.relax.op.prod(x:RelaxExpr, axis:int|List[int] |None= None, keepdims:bool= False) → RelaxExpr
对指定轴上的张量元素进行乘积计算。
- 参数:
- 返回:result : 计算结果。
- 返回类型: relax.Expr。
tvm.relax.op.std(x:RelaxExpr, axis:int|List[int] |None= None, keepdims:bool= False) → RelaxExpr
计算张量元素在给定轴上的标准差。
- 参数:
- 返回:result : 计算结果。
- 返回类型: relax.Expr。
tvm.relax.op.sum(x:RelaxExpr, axis:int|List[int] |None= None, keepdims:bool= False) → RelaxExpr
对指定轴上的张量元素求和。
- 参数:
- 返回:result : 计算结果。
- 返回类型: relax.Expr。
tvm.relax.op.variance(x:RelaxExpr, axis:int|List[int] |None= None, keepdims:bool= False) → RelaxExpr
计算张量元素在给定轴上的方差。
- 参数:
- 返回: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。
输入张量必须具有 float dtype。
tvm.relax.op.acosh(x:RelaxExpr) → RelaxExpr
计算输入数据的逐元素反双曲余弦。
- 参数:
- x (relax.Expr) :输入数据。
- 返回:result :计算结果。
- 返回类型: relax.Expr。
输入张量必须具有 float dtype。
tvm.relax.op.asin(x:RelaxExpr) → RelaxExpr
计算输入数据的逐元素反正弦值。
- 参数:
- x (relax.Expr) :输入数据。
- 返回:result : 计算结果。
- 返回类型: relax.Expr。
输入张量必须具有 float dtype。
tvm.relax.op.asinh(x:RelaxExpr) → RelaxExpr
计算输入数据的逐元素反双曲正弦。
- 参数:
- x (relax.Expr) :输入数据。
- 返回:result :计算结果。
- 返回类型: relax.Expr。
输入张量必须具有 float dtype。
tvm.relax.op.atan(x:RelaxExpr) → RelaxExpr
计算输入数据的逐元素反正切。
- 参数:
- x (relax.Expr) :输入数据。
- 返回:result : 计算结果。
- 返回类型: relax.Expr。
输入张量必须具有 float dtype。
tvm.relax.op.atanh(x:RelaxExpr) → RelaxExpr
计算输入数据的逐元素反双曲正切。
- 参数:
- x (relax.Expr) :输入数据。
- 返回:result :计算结果。
- 返回类型: relax.Expr。
输入张量必须具有 float dtype。
tvm.relax.op.bitwise_not(x:RelaxExpr) → RelaxExpr
计算输入数据的逐位取反。
- 参数:
- 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。
输入张量必须具有 float dtype。
tvm.relax.op.cosh(x:RelaxExpr) → RelaxExpr
计算输入数据的逐元素双曲余弦。
- 参数:
- x (relax.Expr) :输入数据。
- 返回:result :计算结果。
- 返回类型: relax.Expr。
输入张量必须具有 float 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。
输入张量必须具有 float 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。
输入张量必须具有 float dtype。
tvm.relax.op.logical_not(x:RelaxExpr) → RelaxExpr
计算输入数据的逻辑非。
- 参数:
- 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
计算输入数据的逐元素倒数平方根。
- 参数:
- x (relax.Expr) :输入数据
- 返回:result :计算结果。
- 返回类型: relax.Expr。
输入张量必须具有 float dtype。
tvm.relax.op.sigmoid(x:RelaxExpr) → RelaxExpr
计算输入数据的逐元素 Sigmoid 函数。
- 参数:
- x (relax.Expr) :输入数据
- 返回:result :计算结果。
- 返回类型: relax.Expr。
输入张量必须具有 float 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。
输入张量必须具有 float dtype。
tvm.relax.op.sinh(x:RelaxExpr) → RelaxExpr
计算输入数据的逐元素双曲正弦值。
- 参数:
- x (relax.Expr) :输入数据
- 返回:result :计算结果。
- 返回类型: relax.Expr。
输入张量必须具有 float dtype。
tvm.relax.op.sqrt(x:RelaxExpr) → RelaxExpr
计算输入数据的逐元素平方根。
- 参数:
- x (relax.Expr) :输入数据
- 返回:result :计算结果。
- 返回类型: relax.Expr。
输入张量必须具有 float 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。
输入张量必须具有 float dtype。
tvm.relax.op.tanh(x:RelaxExpr) → RelaxExpr
计算输入数据的逐元素 tanh。
- 参数:
- x (relax.Expr) :输入数据
- 返回:result :计算结果。
- 返回类型: relax.Expr。
输入张量必须具有 float dtype。
tvm.relax.op.trunc(x:RelaxExpr) → RelaxExpr
取输入数据的截断值。:param x: 输入数据 :type x: relax.Expr。
- 返回:result :计算结果。
- 返回类型: relax.Expr。
与神经网络相关的运算符。
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)。
- 参数:
- 返回: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 提供了一个整数的元组(高度,宽度),则对于任何输入(NCHW),输出大小为 (N x C x 高度 x 宽度)。
- 参数:
- 返回: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 提供一个整数,则输出大小为(N x C x output_size x output_size x output_size),适用于任何输入(NCDHW)。
如果为 output_size 提供一个整数元组(depth, height, width),则输出大小为(N x C x depth x height x width),适用于任何输入(NCDHW)。
- 参数:
- 返回: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 维张量。
输入张量必须具有 float16 数据类型。
- 参数:
- query (relax.Expr) :运算符的输入查询。输入查询的布局应为 (batch_size, seq_len, num_head, head_dim)。
- key (relax.Expr) :运算符的输入键。输入键的布局应为 (batch_size, seq_len_kv, num_head, head_dim)。
- value (relax.Expr) :运算符的输入值。输入值的布局应为 (batch_size, seq_len_kv, num_head, head_dim_v)。
- bias (Optional[Expr]) :运算符的可选注意力偏置。注意力偏置的布局应为以 seq_len_kv 结尾的 4 维张量,且可广播到 (batch_size, num_head, seq_len, seq_len_kv)。
- scale (Optional[float]) : 应用于注意力分数的缩放值,默认为 1 / sqrt(head_dim)。
- causal_mask (Optional[str]) :可选的因果掩码,即“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 (可选[int]):滑动窗口注意力的窗口大小。
- 返回:result :计算结果。输出布局应为(batch_size, seq_len, num_head, head_dim_v)。
- 返回类型: relax.Expr
tvm.relax.op.nn.attention_bias(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
计算融合多头注意力。
IRModule.script() 将注意力操作转换为注意力偏置,这与 TVMScript 解析器不兼容。该函数使 TVMScript 的打印与 TVMScript 的解析器兼容。
所有输入张量都是 BSNH 布局的 4 维张量。
输入张量必须具有 float16 数据类型。
- 参数:
- query (relax.Expr) :运算符的输入查询。输入查询的布局应为 (batch_size, seq_len, num_head, head_dim)。
- key (relax.Expr) :运算符的输入键。输入键的布局应为 (batch_size, seq_len_kv, num_head, head_dim)。
- value (relax.Expr) :运算符的输入值。输入值的布局应为 (batch_size, seq_len_kv, num_head, head_dim_v)。
- bias (Optional[Expr]) :运算符的可选注意力偏置。注意力偏置的布局应为以 seq_len_kv 结尾的 4 维张量,且可广播到 (batch_size, num_head, seq_len, seq_len_kv)。
- scale (Optional[float]) : 应用于注意力分数的缩放值,默认为 1 / sqrt(head_dim)。
- causal_mask (Optional[str])
:
- 可选的因果掩码,即“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]]
with seq_len = 2, seq_len_kv = 4, mask for ‘TopLeft’:
[[1, 0, 0, 0],
[1, 1, 0, 0]]
“BottomRight”的掩码:
[[1, 1, 1, 0],
[1, 1, 1, 1]]
- window_size (Optional[int]) :滑动窗口注意力的窗口大小。
- 返回: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]) :查询序列长度的累积和,以 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]) :应用于注意力分数的缩放值,默认为 1 / sqrt(head_dim)。
- causal_mask (Optional[str]) :
- 可选的因果掩码,即“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]]
with seq_len = 2, seq_len_kv = 4, mask for ‘TopLeft’:
[[1, 0, 0, 0],
[1, 1, 0, 0]]
“BottomRight”的掩码:
[[1, 1, 1, 0],
[1, 1, 1, 1]]
- window_size (Optional[int]) : 滑动窗口注意力的窗口大小。
- 返回: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= True, layout:str= 'NCW', out_layout:str|None= None) → RelaxExpr
1D 平均池化运算符。
此运算符接收数据作为输入,并通过 stride 定义的步长,使用 pool_size 大小的窗口进行 1D 平均值计算。
在默认情况下,当 data_layout 为 NCW 时,一个数据张量 shape 为(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]**]) :池化的填充。必须具有长度为 1 或 2。
- dilation (Union[int,Tuple[int]**]) :池化的膨胀。必须具有长度为 1。
- ceil_mode (bool) : 一个布尔值,指示是否使用 ceil 或 floor 来计算输出形状。使用 ceil 时,输入张量的每个元素都将被滑动窗口覆盖。
- count_include_pad (bool,optional) :是否将填充包括在内以计算平均值。
- layout (str) :输入的布局。
- out_layout (Optional[str]) :输出的布局。如果未指定,则与 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 平均池化运算符。
该运算符以数据为输入,通过指定步长 stride,在 pool_size 大小的窗口内进行 2D 平均值计算。
在默认情况下,当数据布局为 NCHW 时,一个形状为(batch_size, in_channels, height, width)的数据张量,将按照以下规则产生输出张量:
在计算之前对数据进行填充。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, …**]]) :池化操作的填充。必须具有长度为 1、2 或 4。
- dilation (Union[int,Tuple[int, int]**]) : 池化操作的膨胀率。必须具有长度为 1 或 2。
- ceil_mode (bool) :一个布尔值,指示是否使用 ceil 或 floor 来计算输出形状。使用 ceil 时,输入张量的每个元素都将被滑动窗口覆盖。
- count_include_pad (bool,optional) :是否将填充包括在内以计算平均值。
- layout (str) :输入的布局。
- out_layout (Optional[str]) : 输出的布局。如果未指定,则与 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= True, layout:str= 'NCDHW', out_layout:str|None= None) → RelaxExpr
2D 平均池化运算符。
这个运算符以数据为输入,通过 stride 定义的步长,使用 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) :一个布尔值,指示是否使用 ceil 或 floor 来计算输出形状。使用 ceil 时,输入张量的每个元素都将被滑动窗口覆盖。
- count_include_pad (bool,optional) :是否将填充包括在内以计算平均值。
- layout (str) :输入的布局。
- out_layout (Optional[str]) :输出的布局。如果未指定,则与 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, training:bool= True) → RelaxExpr
Batch normalization 层(Ioffe 和 Szegedy,2014)。
对每个批次进行输入规范化,即应用一个保持激活均值接近 0 和激活标准差接近 1 的变换。
均值和方差都通过将输入视为向量来返回标量。
然后计算归一化输出,其形状与输入相同,如下所示:
假设输入在轴 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 将通道轴设置为输入形状中的最后一项。
这个运算符有两种模式:
- 训练模式。
- 使用从当前批次计算出的均值和方差进行归一化。
- 更新并返回运行均值和运行方差。
- 推理模式。
- 使用 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) :指示是否将伽马缩放相乘。
- momentum (float) :用于 moving_mean 和 moving_var 更新的值。
- training (bool) : relax batch_norm 处于训练模式。要将其转换为推理模式,可以使用 DecomposeOpsForInference。
- 返回: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| dtype |None= None) → RelaxExpr
1D 卷积。
该运算符将权重作为 1D 卷积核,并将其与数据卷积以产生输出。
在默认情况下,其中数据布局为 NCW,核布局为 OIW,conv1d 接收一个数据张量,其形状为(batch_size, in_channels, width),以及一个权重张量,其形状为(channels, in_channels, kernel_w),其中 kernel_w 是 W 核维度的长度,以产生一个输出张量,其规则如下:
在计算之前,数据张量和权重张量分别应用填充和膨胀。该运算符接受数据布局规范。从语义上讲,该运算符将布局转换为规范布局(数据为 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]) :输出的布局。如果未指定,则与 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| dtype |None= None) → RelaxExpr
1D 转置卷积运算符。
这个运算符可以看作是 conv1d 的梯度运算符。
输出形状可以在 data_layout == “NCW”和 kernel_layout == “IOW”的简单情况下解释。假设数据形状为(N, in_channel, in_w),权重形状为(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, …]],optional) :用于区分输出形状。
- dilation (Union[int,Tuple[int]**]) :指定用于扩张卷积的扩张率。必须具有长度为 1。
- groups (int) :分组卷积将输入分成多少组。输入和输出通道数应该能被组数整除。
- data_layout (str) :输入的布局。
- kernel_layout (str) :权重的布局。
- out_layout (Optional[str]) : 输出的布局。如果未指定,则与 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| dtype |None= None) → RelaxExpr
二维卷积。
该运算符将权重作为卷积核,并将其与数据卷积以产生输出。
在默认情况下,其中数据布局为 NCHW,核布局为 OIHW,conv2d 接收一个形状为(batch_size,in_channels,height,width)的数据张量,以及一个形状为(channels,in_channels,kernel_h,kernel_w)的权重张量,其中 kernel_h 和 kernel_w 是 H 和 W 核维度的长度,以产生一个输出张量,其规则如下:
在计算之前,对数据和权重分别应用填充和膨胀。该运算符接受数据布局规范。从语义上讲,运算符将布局转换为规范布局(数据为 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]) : 输出的布局。如果未指定,则与 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| dtype |None= None) → RelaxExpr
二维转置卷积运算符。
这个运算符设计为 conv2d 的梯度运算符。这意味着,如果
out = conv2d(data, weight, strides, padding, dilation),
那么相对于 data 的梯度可以按如下方式计算:
data_grad = conv2d_transpose(out_grad, weight, strides, padding, output_padding, dilation),output_padding 是一个用于确定输出形状的参数。
当 data_layout == “NCHW”且 kernel_layout == “IOHW”时,输出形状可以简单解释。假设数据形状为(N, in_channel, in_h, in_w),权重形状为(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, …]]**,optional) :用于区分输出形状。
- dilation (Union*[int,Tuple[int, int]**]) :指定用于扩张卷积的扩张率。必须具有长度为 1 或 2。
- groups (int) :分组卷积将输入分成多少组。输入和输出通道数应该能被组数整除。
- data_layout (str) : 输入的布局。
- kernel_layout (str) : 权重的布局。
- out_layout (Optional[str]) : 输出的布局。如果未指定,则与 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| dtype |None= None) → RelaxExpr
3D 卷积。
该运算符将权重作为卷积核,将其与数据卷积以产生输出。
在默认情况下,当 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 卷积核维度的长度,以产生一个输出张量,其规则如下:
在计算之前,数据会应用填充,权重会应用膨胀。该运算符接受数据布局规范。从语义上讲,运算符会将布局转换为规范布局(数据为 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]) :输出的布局。如果未指定,则与 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
predictions 和 labels 之间的 logits 交叉熵。
predictions 和 labels 的形状必须相同。当 ndim >= 2 时,第一个维度被视为批次大小 N。在这种情况下,计算结果将除以 N 以进行均值归约。
- 参数:
- predictions (relax.Expr) :预测值。
- labels (relax.Expr) :标签(真实值)。
- 返回:result :计算结果。
- 返回类型: relax.Expr。
tvm.relax.op.nn.dropout(data:RelaxExpr, rate:float= 0.5) → RelaxExpr
对输入张量应用 dropout 操作。
在训练期间,输入的每个元素以概率 p
被设置为 0。整个数组被缩放为 1/(1-p)
以保持输入的期望和不变。
- 参数:
- data (relax.Expr) :运算符的输入数据。
- rate (float) :元素重置为 0 的概率。
- 返回:result :dropout 的结果,是一个包含两个张量的元组。第一个是原始张量,第二个是掩码张量(未丢弃元素处为 1.0,丢弃处为 0.0)。
- 返回类型: relax.Expr。
tvm.relax.op.nn.gelu(data:RelaxExpr) → RelaxExpr
高斯误差线性单元函数
其中 erf 是高斯误差函数。
- 参数:
- data (relax.Expr) *:*输入数据。
- 返回:result :计算结果。
- 返回类型: relax.Expr。
输入张量必须具有 float dtype。
tvm.relax.op.nn.gelu_tanh(data:RelaxExpr) → RelaxExpr
高斯误差线性单元函数,使用 tanh 近似
- 参数:
- data (relax.Expr) :输入数据。
- 返回:result :计算结果。
- 返回类型: relax.Expr。
输入张量必须具有 float 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
分组归一化(Yuxin Wu 等人,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) :指示是否将伽马缩放相乘。
- 返回:result :计算结果。
- 返回类型: relax.Expr。
tvm.relax.op.nn.instance_norm(data:RelaxExpr, gamma:RelaxExpr, beta:RelaxExpr, channel_axis:int, axes:List[int], epsilon:float= 1e-05, center:bool= True, scale:bool= True) → RelaxExpr
参数:
- data (relax.Expr) : 应用 instance_norm 的输入。
- gamma (relax.Expr) : gamma 缩放因子。
- beta (relax.Expr) :beta 偏移因子。
- axes (Union*[int,*** List*[int]]*) :正在应用归一化的轴。
- epsilon (float) :添加到方差中的小浮点数,以避免除以零。
- center (bool) :指示是否将 beta 偏移添加到归一化张量中。
- scale (bool) :指示是否将伽马缩放相乘。
- 返回: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
层归一化(Lei Ba 等人,2016 年)。将层归一化应用于 n 维输入数组。该运算符接收一个 n 维输入数组,并使用指定的轴对输入进行归一化:
与批量归一化不同,均值和方差是沿着通道维度计算的。
假设输入在轴 1 上的大小为 k,那么 gamma 和 beta 的形状都是(k,)。
这个运算符可以在推理时被优化掉。
- 参数:
- 返回:result : 计算结果。
- 返回类型: relax.Expr。
tvm.relax.op.nn.leakyrelu(data:RelaxExpr, alpha:float= 0.01) → RelaxExpr
修正线性单元。
- 参数:
- data (relax.Expr) :输入数据。
- alpha (float) :控制负斜率的角,用于负输入。默认值为 0.01
- 返回:result :计算结果。
- 返回类型: relax.Expr。
tvm.relax.op.nn.log_softmax(data:RelaxExpr, axis:int= -1) → RelaxExpr
计算 log softmax。
这个运算符可以在推理时被优化掉。
- 参数:
- 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
1D 最大池化运算符。
此运算符以数据为输入,通过 stride 定义的步长,在 pool_size 大小的窗口内进行 1D 最大值计算。
在默认情况下,当 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) :一个布尔值,指示是否使用 ceil 或 floor 来计算输出形状。使用 ceil 时,输入张量的每个元素都将被滑动窗口覆盖。
- count_include_pad (bool,optional) : 是否将填充包括在内以计算平均值。
- layout (str) *:*输入的布局。
- out_layout (Optional[str]) :输出的布局。如果未指定,则与 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
2D 最大池化运算符。
该运算符接收数据作为输入,并通过 stride 定义的步长,在 pool_size 大小的窗口内进行 2D 最大值计算。
在默认情况下,当数据布局为 NCHW 时,一个形状为(batch_size, in_channels, height, width)的数据张量,将按照以下规则产生输出张量:
形状为(b, c, h, w)的数据和池化大小(kh, kw)
在计算之前对数据进行填充。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, …**]]) : 池化操作的填充。必须具有长度为 1、2 或 4。
- dilation (Union*[int,*** Tuple[int, int]**]) : 池化操作的膨胀率。必须具有长度为 1 或 2。
- ceil_mode (bool) : 一个布尔值,指示是否使用 ceil 或 floor 来计算输出形状。使用 ceil 时,输入张量的每个元素都将被滑动窗口覆盖。
- count_include_pad (bool,optional) :是否将填充包括在内以计算平均值。
- layout (str) :输入的布局。
- out_layout (Optional[str]) :输出的布局。如果未指定,则与 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
3D 最大池化运算符。
该运算符接收数据作为输入,并通过 stride 定义的步长,使用 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]**]) :池化窗口的大小。必须具有 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) :一个布尔值,指示是否使用 ceil 或 floor 来计算输出形状。使用 ceil 时,输入张量的每个元素都将被滑动窗口覆盖。
- count_include_pad (bool,optional) :是否将填充包括在内以计算平均值。
- layout (str) :输入的布局。
- out_layout (Optional[str]) :输出的布局。如果未指定,则与 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, i1, i2, …, i_k] = -p * w, where - p = predictions[n, t, i1, i2, i_k], - t = targets[n, i1, i2, …, i_k], - w = weights[t] if t != ignore_index else 0。
result = reduction(output)。
- 参数:
- 返回:result :计算结果。
- 返回类型: relax.Expr。
tvm.relax.op.nn.pad(data:RelaxExpr, pad_width:List[int] |Tuple[int, …], pad_mode:str|None= 'constant', pad_value:float|None= 0.0)
这个运算符接收一个张量,并使用指定的值按照指定的宽度对每个轴进行填充。
- 参数:
- data (relax.Expr) :运算符的输入数据
- pad_width (Union*[List[int**],Tuple[int,…]***],required) :每个轴边缘填充的值数量,格式为 ((before_1, after_1), …, (before_N, after_N))
- pad_mode (Optional[str]) :‘constant’, ‘reflect’, ‘replicate’, ‘circular’ ‘constant’ 使用常数值填充 pad_value ‘reflect’ 通过镜像值(不包括边缘)进行填充 ‘replicate’ 通过重复边缘值进行填充。 ‘circular’ 通过从另一侧循环值进行填充。默认为 ‘constant’
- pad_value (Optional*[Union[float,* Expr*]]*) :用于填充的值。默认为 0。
- 返回:result :计算结果。
- 返回类型: relax.Expr。
tvm.relax.op.nn.pixel_shuffle(data:RelaxExpr, upscale_factor:int)
像素重排运算符。
该运算符对输入张量执行像素重排操作,常用于图像超分辨率任务中的高效亚像素卷积。它将形状为 (N, C × r^2, H, W) 的张量中的元素重新排列为形状为 (N, C, H × r, W × r) 的张量,其中 r 是放大因子。
- 参数:
- data (relax.Expr) :像素重排运算符的输入张量。它必须有 4 个维度,格式为 (N, C 运算符* r^2, H, W),其中 r 是放大因子。
- upscale_factor (int) :放大因子 r。它决定了输入张量的空间分辨率(高度和宽度)增加的量。
- 返回:result : 变换后的张量,形状为 (N, C, H 运算符* r, W 运算符* r)。
- 返回类型: relax.Expr。
示例
如果输入张量的形状为 (1, 8, 10, 15) 且上采样因子为 2,则结果张量的形状将为 (1, 2, 20, 30)。
tvm.relax.op.nn.prelu(data:RelaxExpr, alpha:RelaxExpr, axis:int= 1) → RelaxExpr
参数化整流线性单元(PReLU)。
- 参数:
- 返回:result :计算结果。
- 返回类型: relax.Expr。
tvm.relax.op.nn.relu(data:RelaxExpr) → RelaxExpr
修正线性单元。
- 参数:
- data (relax.Expr) : 输入数据。
- 返回:result :计算结果。
- 返回类型: relax.Expr。
tvm.relax.op.nn.relu6(data:RelaxExpr) → RelaxExpr
ReLU6 激活函数。
- 参数:
- 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 维输入数组,并使用指定轴对输入进行归一化:
- 参数:
- 返回:result :计算结果。
- 返回类型: relax.Expr。
tvm.relax.op.nn.selu(data:RelaxExpr) → RelaxExpr
缩放指数线性单元(SELU)。
\begin{split}\text{SELU}(x) = \lambda \begin{cases} x & \text{if } x > 0 \ \alpha (e^x - 1) & \text{if } x \leq 0 \end{cases}\end{split}
其中 λ≈1.0507 和 α≈1.6733 。
- 参数:
- data (relax.Expr) :输入数据。
- 返回:result : 计算结果。
- 返回类型: relax.Expr
tvm.relax.op.nn.silu(data:RelaxExpr) → RelaxExpr
Sigmoid 线性单元函数。
- 参数:
- data (relax.Expr) : 输入数据。
- 返回:result :计算结果。
- 返回类型: relax.Expr。
输入张量必须具有 float dtype。
tvm.relax.op.nn.softmax(data:RelaxExpr, axis:int= -1) → RelaxExpr
计算 softmax。
- 参数:
- data (relax.Expr) :运算符的输入数据。
- axis (int) : 计算 softmax 时沿其求和的轴。如果未指定,默认为输入张量的最后一个轴。支持负索引。
- 返回:result :计算结果。
- 返回类型: relax.Expr。
输入张量必须具有 float dtype。
tvm.relax.op.nn.softplus(data:RelaxExpr, beta:float= 1.0, threshold:float= 20.0) → RelaxExpr
Softplus 激活函数。
- 参数:
- 返回:result :计算结果。
- 返回类型: relax.Expr。
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 的张量。
- 参数:
- 返回:result :一个 relax relax.Call,它获取已分配的张量。
- 返回类型:relax.Call。
tvm.relax.op.builtin.stop_lift_params(x:RelaxExpr) → RelaxExpr
一个指示输入张量的消费者不应被提升到 transform_params 函数的标志。
- 参数:
- x (relax.Expr) :输入数据。
- 返回:result : 结果张量,与输入张量相同。
- 返回类型: relax.Expr。
与 CCL 相关的运算符。
tvm.relax.op.ccl.allgather(x, num_workers:int, in_group:bool= True)
AllGather 运算符。
- 参数:
- 返回:result :allgather 的结果。
- 返回类型: relax.Expr。
tvm.relax.op.ccl.allreduce(x, op_type:str= 'sum', in_group:bool= True)
Allreduce 运算符。
- 参数:
- 返回:result :allreduce 的所有结果。
- 返回类型: relax.Expr。
tvm.relax.op.ccl.broadcast_from_worker0(x:RelaxExpr) → RelaxExpr
从 worker-0 广播数据到所有其他 worker。
- 参数:x (relax.Expr) : 要广播的张量。
- 返回:result : 已广播到所有其他 worker 的相同张量。
- 返回类型: relax.Expr。
tvm.relax.op.ccl.scatter_from_worker0(x:RelaxExpr, num_workers:int, axis:int= 0) → RelaxExpr
执行从 worker-0 开始的 scatter 操作,将给定的 buffer 分成相等的部分。
- 参数:
- 返回:result : 由不同工作进程接收的 Chunked Tensor。
- 返回类型: 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 应该是一个工作本地函数,实际上在每个工作节点上执行,而不是未分割的函数。这个运算符的输出是 DTensor 或 DTensor 的元组。
- 参数:
- gvar (GlobalVar) :指向 tir PrimFunc 的 GlobalVar。
- args (Expr) :输入参数。
- out_sinfo (Union*[DTensorStructInfo,* List*[DTensorStructInfo**]]*) :调用_tir 的输出结构信息。它应该是一个 DTensorStructInfo 或一个 DTensorStructInfo 的列表。每一个表示一个返回张量的结构信息。
- tir_vars (Optional*[Union[ShapeExpr,* Tuple[PrimExpr],List[PrimExpr]**]]) :表示调用 func 时需要解包的整数元组的 ShapeExpr。如果未使用则为 null。
- 返回: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 操作的专用版本。
- 参数:
- 返回: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 :对数据的梯度。
- 返回类型: 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 相同。返回相对于数据的梯度。
- 参数:output_grad (relax.Expr) :对 max_pool2d 结果的梯度。
- 返回:result : 对数据的梯度。
- 返回类型: 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 相同。返回对预测的梯度。
- 参数:output_grad (relax.Expr) :对 nll_loss 结果的梯度。
- 返回:result:对预测的梯度。
- 返回类型: 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.75, cubic_exclude:int= 0, extrapolation_value:float= 0.0, out_dtype:str| dtype |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,其必须具有维度 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) :输入的布局。
- method (str) : 使用的缩放方法 [最近邻, 线性, 三次插值]。
- coordinate_transformation_mode (str) :描述如何将调整大小后的张量中的坐标转换为原始张量中的坐标。定义可以在 topi/image/resize.py 中找到。[半像素, 对齐角点, 非对称, pytorch_half_pixel, tf_half_pixel_for_nn, 和 tf_crop_and_resize]。
- rounding_method (str) :指示在最近邻方法中如何找到“最近”的像素 [四舍五入, 向下取整, 向上取整
- 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 来分配具有特定大小、虚拟设备索引、存储范围和 dtype 的存储。
- 参数:
- 返回: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 来在指定的存储上从给定偏移量开始分配一个张量。
- 参数:
- 返回: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。
tvm.relax.op.memory.view(data:RelaxExpr, shape:Sequence[int|PrimExpr] |RelaxExpr|None= None, dtype:RelaxExpr|None= None, relative_byte_offset:RelaxExpr|None= None) → RelaxExpr
提供一个现有张量的视图。
视图的形状可能不同,数据类型可能不同,并且可能相对于源数组从偏移量开始。
无论使用这些选项的哪种组合,视图都不会访问通过输入数据数组无法访问的内存。即使数据数组本身是共享后备数组的视图,此限制也适用。
- 参数:
- 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 的偏移量,相对于数据的字节偏移量。如果为 None,视图的偏移量与数据的偏移量相同。
- 返回:result :张量视图。
- 返回类型: relax.Expr。
tvm.relax.op.memory.ensure_zero_offset(data:RelaxExpr) → RelaxExpr
确保张量具有 elem_offset == 0。如有必要,将进行复制。
- 参数:
- data (relax.Expr) ***:***输入张量。
- Results。
- -------
- result (relax.Expr) *:*元素偏移量等于 0 的张量。
tvm.relax.op.op_attrs
用于 Relax 操作的属性节点。
class tvm.relax.op.op_attrs.CallTIRWithGradAttrs
call_tir_with_grad 操作使用的属性。
propertyte_grad_kwargs
传递给 te gradient 函数的关键字参数。
propertyte_grad_name
与这个 call_tir_with_grad 节点关联的 te 梯度函数的名称。
class tvm.relax.op.op_attrs.InitAttrs
用于 full/full_like、ones/ones_like 和 zeros/zeros_like 运算符的属性。
propertydtype
创建的张量的数据类型。
class tvm.relax.op.op_attrs.TriluAttrs
tril 和 triu 运算符中使用的属性。
propertyk
主对角线上方或下方要排除或包含的对角线数量。
class tvm.relax.op.op_attrs.AstypeAttrs
astype 运算符中使用的属性。
propertydtype
目标数据类型。
class tvm.relax.op.op_attrs.TakeAttrs
take 运算符中使用的属性。
propertyaxis
选择值的轴。
propertymode
处理越界索引的模式。
class tvm.relax.op.op_attrs.StridedSliceAttrs
strided_slice 运算符使用的属性。
propertyassume_inbound
是否假设索引在有效范围内。如果设置为 false,超出范围的索引将被裁剪到边界。
class tvm.relax.op.op_attrs.MatmulAttrs
matmul 运算符的属性。
propertyout_dtype
输出张量的数据类型。
class tvm.relax.op.op_attrs.Conv2DAttrs
nn.conv2d 的属性
propertydata_layout
输入数据的维度顺序。可以是‘NCHW’、‘NHWC’等。‘N’、‘C’、‘H’、‘W’分别代表批次、通道、高度和宽度维度。卷积应用于‘H’和‘W’维度。
propertydilation
指定用于扩张卷积的扩张率。
propertygroups
分组卷积将输入分成多少组。输入和输出通道的数量应该能被组数整除。
propertykernel_layout
权重的维度排序。可以是‘OIHW’、‘OIHW16o16i’等。‘O’、‘I’、‘H’、‘W’分别代表 num_filter、input_channel、height 和 width 维度。
propertyout_dtype
输出数据类型,在混合精度设置下设为显式类型。
propertyout_layout
输出维度的排序。可以是‘NCHW’、‘NHWC’等。'N'、'C'、'H'、'W'分别代表批次、通道、高度和宽度维度。默认与输入布局相同。
propertypadding
底部和右侧将使用与顶部和左侧相同的填充值四个 int : 填充宽度按 (顶部, 左侧, 底部, 右侧) 的顺序排列。
Type : 所有边使用相同的填充两个整数。
propertystrides
指定卷积的步长。
class tvm.relax.op.op_attrs.Conv3DAttrs
nn.conv3d 的属性。
propertydata_layout
输入数据的维度顺序。可以是 ‘NCDHW’、‘NDHWC’ 等。'N'、'C'、'D'、'H'、'W' 分别代表批次、通道、深度、高度和宽度维度。卷积在 ‘D’、‘H’ 和 ‘W’ 维度上应用。
propertydilation
指定用于扩张卷积的扩张率。
propertygroups
分组卷积将输入分成多少组。输入和输出通道的数量应该能被组数整除。
propertykernel_layout
权重的维度排序。可以是“OIDHW”、“OIDHW16o16i”等。“O”、“I”、“D”、“H”、“W”分别代表 num_filter、input_channel、depth、height 和 width 维度。
propertyout_dtype
输出数据类型,在混合精度设置下设为显式类型。
propertyout_layout
输出维度的排序。可以是‘NCDHW’、‘NDHWC’等。'N'、'C'、'D'、'H'、'W'分别代表批次、通道、深度、高度和宽度维度。默认与输入布局相同。
propertypadding
底部和右侧将使用与顶部和左侧相同的填充值六整数:填充宽度按顺序为(前,后,上,左,下,右)。
**Type :**如果填充非零,则输入隐式零填充 Padding 支持对称和非对称作为一整数。
**Type :**所有边使用相同的填充两个整数。
propertystrides
指定卷积的步长。
class tvm.relax.op.op_attrs.Conv2DTransposeAttrs
nn.conv2d_transpose 的属性。
propertydata_layout
输入数据的维度顺序。可以是‘NCHW’、‘NHWC’等。‘N’、‘C’、‘H’、‘W’分别代表批次、通道、高度和宽度维度。卷积应用于‘H’和‘W’维度。
propertydilation
指定用于扩张卷积的扩张率。
propertygroups
分组卷积将输入分成多少组。输入和输出通道的数量应该能被组数整除。
propertykernel_layout
权重的维度排序。可以是‘OIHW’、‘OIHW16o16i’等。‘O’、‘I’、‘H’、‘W’分别代表 num_filter、input_channel、height 和 width 维度。
propertyout_dtype
输出数据类型,在混合精度设置下设为显式类型。
propertyout_layout
输出维度的排序。可以是‘NCHW’、‘NHWC’等。'N'、'C'、'H'、'W'分别代表批次、通道、高度和宽度维度。默认与输入布局相同。
propertyoutput_padding
用于消除输出形状的歧义。
propertypadding
底部和右侧将使用与顶部和左侧相同的填充值四个 int : 填充宽度按 (顶部, 左侧, 底部, 右侧) 的顺序排列。
**Type :**如果填充非零,则输入隐式零填充 Padding 支持对称和非对称作为一整数。
**Type :**所有边使用相同的填充两个整数。
propertystrides
指定卷积的步长。
class tvm.relax.op.op_attrs.Pool2DAttrs
nn.max_pool2d 的属性。
propertyceil_mode
一个布尔值,表示是否使用向上取整或向下取整来计算输出形状。使用向上取整时,输入张量的每个元素都将被滑动窗口覆盖。
propertycount_include_pad
当为真时,将包含填充以计算平均值。
propertydilation
指定卷积的膨胀率。
propertylayout
输入数据的维度顺序。可以是‘NCHW’、‘NHWC’等。’N’、‘C’、‘H’、‘W’分别代表批次、通道、高度和宽度维度。池化操作应用于‘H’和‘W’维度。
propertyout_layout
输出数据的维度顺序。可以是‘NCHW’、‘NHWC’等。’N’、‘C’、‘H’、‘W’分别代表批次、通道、高度和宽度维度。池化操作应用于‘H’和‘W’维度。
propertypadding
底部和右侧将使用与顶部和左侧相同的填充值四个 int : 填充宽度按 (顶部, 左侧, 底部, 右侧) 的顺序排列。
Type : 如果填充非零,则输入隐式零填充 Padding 支持对称和非对称作为一整数。
Type : 所有边使用相同的填充两个整数。
propertypool_size
池化窗口的大小。
propertystrides
指定卷积的步长。
class tvm.relax.op.op_attrs.AdaptivePool2DAttrs
2d 自适应池化运算符的属性。
propertylayout
输入数据的维度顺序。可以是‘NCHW’、‘NHWC’等。’N’、‘C’、‘H’、‘W’分别代表批次、通道、高度和宽度维度。池化操作应用于‘H’和‘W’维度。
propertyout_layout
输出数据的维度顺序。可以是‘NCHW’、‘NHWC’等。’N’、‘C’、‘H’、‘W’分别代表批次、通道、高度和宽度维度。池化操作应用于‘H’和‘W’维度。
propertyoutput_size
输出高度和宽度。
class tvm.relax.op.op_attrs.SoftmaxAttrs
nn.softmax 运算符的属性。
propertyaxis
计算 softmax 时要对其求和的轴。
class tvm.relax.op.op_attrs.BatchNormAttrs
批量归一化运算符使用的属性
propertyaxis
归一化应用的轴。
propertycenter
指示是否将 beta 偏移量添加到归一化张量。
propertyepsilon
为方差添加的小浮点数以避免除以零。
propertymomentum
用于移动平均均值和移动方差更新的值。
propertyscale
指示是否将 gamma 尺度相乘。
propertytraining
我们是否在训练(即不在评估模式下)。
class tvm.relax.op.op_attrs.LayerNormAttrs
层归一化运算符使用的属性。
propertyaxes
归一化应用的轴。
propertycenter
指示是否将 beta 偏移量加到归一化张量上。
propertyepsilon
一个小的浮点数加到方差上以避免除以零。
propertyscale
指示是否将伽马尺度相乘。
class tvm.relax.op.op_attrs.InstanceNormAttrs
instance_norm 操作使用的属性
propertyaxes
归一化应用的轴。
propertycenter
指示是否将 beta 偏移量添加到归一化张量。
propertychannel_axis
表示通道的轴。
propertyepsilon
为方差添加的小浮点数以避免除以零。
propertyscale
指示是否将 gamma 缩放相乘。
class tvm.relax.op.op_attrs.DropoutAttrs
dropout 运算符的属性。
propertyrate
训练期间被丢弃的输入比例。
class tvm.relax.op.op_attrs.StatisticalAttrs
统计运算符中使用的属性。
propertyaxis
执行归约的轴或轴。
propertykeepdims
如果设置为 True,缩减的轴会以大小为 1 的维度保留在结果中。
class tvm.relax.op.op_attrs.ConcatAttrs
concat 运算符的属性。
propertyaxis
输入数组连接的轴。应在 [-ndim, ndim) 范围内。
class tvm.relax.op.op_attrs.ExpandDimsAttrs
expand_dims 运算符的属性。
propertyaxis
输入数组扩展的轴。所有值都必须位于范围 [-data.ndim - 1, data.ndim] 内,遵循负索引的约定。
class tvm.relax.op.op_attrs.PermuteDimsAttrs
permute_dims 运算符的属性。
propertyaxes
目标轴顺序,如果未指定则逆序。
class tvm.relax.op.op_attrs.SortAttrs
排序运算符的属性
propertyaxis
计算排序的轴。默认使用最后一个轴。
propertydescending
是否按降序排序。如果未指定,则默认为升序。
class tvm.relax.op.op_attrs.ArgsortAttrs
argsort 运算符的属性。
propertyaxis
argsort 计算的轴。默认使用最后一个轴。
propertydescending
是否以降序 argsort。如果未指定,默认为升序。
propertydtype
输出索引的数据类型。
class tvm.relax.op.op_attrs.SplitAttrs
split 运算符使用的属性。
propertyaxis
要拆分的轴。
propertyindices_or_sections
索引输入数组或分割部分的数目。
class tvm.relax.op.op_attrs.SqueezeAttrs
squeeze 运算符的属性。
propertyaxis
在输入张量中要压缩的轴。如果 axis = None,则压缩维度 1 的所有轴;否则,压缩 axes 中指定的维度。如果轴的维度不是 1,则会出错。
class tvm.relax.op.op_attrs.StackAttrs
concat 运算符的属性。
propertyaxis
用于堆叠输入张量的轴。该轴将被插入到输出的这个位置,因此它必须在范围 [-ndim-1, ndim] 内,其中 ndim 是输入张量的维度数量。
class tvm.relax.op.op_attrs.IndexPutAttrs
index_put 运算符的属性。
propertyaccumulate
是否累积(相加)值而不是替换。如果为真,执行 tensor[indices] += values,否则执行 tensor[indices] = values。
class tvm.relax.op.op_attrs.LayoutTransformAttrs
布局转换操作中使用的属性。
propertyaxis_separators
生成扁平输出轴时输入轴之间的分隔符。
propertyindex_map
要应用的空间变换。
propertyinput_axis_separators
重新生成输出时的轴之间的分隔符。
propertypad_value
如果布局转换会导致隐式填充,则用于填充的特定值。如果未指定,编译器可以自由选择任何值。
class tvm.relax.op.op_attrs.Resize2DAttrs
图像 resize2d 操作中使用的属性。
propertycoordinate_transformation_mode
描述了如何将调整大小后的张量的坐标转换为原始张量的坐标。详细内容请参阅 ONNX Resize 算子规范。可选的选项有 half_pixel、align_corners 和 asymmetric。
propertycubic_alpha
双三次插值的三次样条系数。
propertycubic_exclude
用于在双三次插值过程中排除图像外部的标志。
propertyextrapolation_value
当 roi 在图像外部时返回的值。
propertylayout
输入数据的维度顺序。可以是‘NCHW’、‘NHWC’等。‘N’、‘C’、‘H’、‘W’分别代表批次、通道、高度和宽度维度。Resize 操作应用于‘H’和‘W’维度。
propertymethod
指定用于缩放的模式。nearest_neighbor - 最近邻 linear - 双线性插值 cubic - 双三次插值。
propertyout_dtype
输出张量的 dtype。如果未指定,输出将具有与输入相同的 dtype(如果未指定)。
propertyroi
感兴趣区域(Region of Interest)用于坐标变换模式‘tf_crop_and_resize’。
propertyrounding_method
指示在最近邻方法中如何找到“最近”的像素。可用选项有 round、floor 和 ceil。
class tvm.relax.op.op_attrs.ArgmaxArgminAttrs
argmax/argmin 运算符的属性。
propertyaxis
执行 argmin/argmax 的轴。
propertykeepdims
如果设置为 True,则缩减的轴会以大小为 1 的维度保留在结果中。
class tvm.relax.op.op_attrs.RepeatAttrs
repeat 运算符的属性。
propertyaxis
沿哪个轴重复值。负数从后向前计数。默认情况下,使用展平的输入数组,并返回一个展平的输出数组。
propertyrepeats
重复次数。
class tvm.relax.op.op_attrs.TileAttrs
分块运算符的属性。
propertyrepeats
每个轴上数据重复的次数。
class tvm.relax.op.op_attrs.ScanopAttrs
扫描运算符的属性。
propertyaxis
执行扫描计算沿的轴。默认值(None)是对展平数组进行计算。
propertydtype
输出数据类型。如果未指定 dtype,则默认为输入数据的 dtype。
propertyexclusive
第一个元素不包括。
class tvm.relax.op.op_attrs.TopKAttrs
topk 运算符的属性。
propertyaxis
沿此轴对输入张量进行排序。
propertydtype
输出索引的数据类型。
propertyk
选择顶部元素的数量。
propertylargest
是否返回最大或最小的元素。默认情况下,返回最大的 k 个元素。
propertyret_type
返回类型 [both, values, indices]。both - 返回 top k 数据和索引.values - 仅返回 top k 数据.indices - 仅返回 top k 索引。
class tvm.relax.op.op_attrs.EinsumAttrs
einsum 运算符的属性。
propertysubscripts
einsum 表达式字符串。
class tvm.relax.op.op_attrs.FlipAttrs
翻转运算符的属性。
propertyaxis
翻转的轴。
classtvm.relax.op.op_attrs.PadAttrs
pad 算子中使用的属性。
propertypad_mode
使用的填充类型。“constant”用 constant_value 进行填充,“edge”使用输入数组的边缘值进行填充,“reflect”通过相对于边缘反射值进行填充。
propertypad_value
用于填充填充区域的值。
propertypad_width
每个轴边缘填充的值的数量,格式为(before_1, after_1, …, before_N, after_N)。
classtvm.relax.op.op_attrs.MultinomialFromUniformAttrs
multinomial_from_uniform 算子的属性。
propertydtype
输出索引的数据类型。
classtvm.relax.op.op_attrs.CallInplacePackedAttrs
call_inplace_packed 操作符使用的属性。
classtvm.relax.op.op_attrs.CallTIRInplaceAttrs
call_tir_inplace 操作符使用的属性。
classtvm.relax.op.op_attrs.ToVDeviceAttrs
to_vdevice 操作符使用的属性。
propertydst_vdevice
数据被复制到的目标设备。
classtvm.relax.op.op_attrs.HintOnDeviceAttrs
用于 hint_on_device 操作符的属性。
propertydev_id
设备 ID。
propertydev_type
数据预期执行设备类型。
classtvm.relax.op.op_attrs.ScatterCollectiveAttrs
scatter 集体操作符中使用的属性。
propertyaxis
要散布的张量的轴。张量将沿此轴分块。
propertynum_workers
工作线程的数量,也是给定缓冲区应分块成部分的数量。
classtvm.relax.op.op_attrs.AttentionAttrs
用于注意力算子的属性。
propertycausal_mask
因果掩码的类型,即“TopLeft”和“BottomRight”。
propertyscale
在 softmax 之前应用的定制缩放。默认值为 1 / sqrt(head_dim)。
propertywindow_size
滑动窗口注意力的窗口大小。
classtvm.relax.op.op_attrs.Conv1DAttrs
nn.conv1d 的属性。
propertydata_layout
输入数据的维度排序。可以是 ‘NCW’、‘NWC’ 等。'N'、'C'、'W' 分别代表批次、通道、宽度。卷积应用于 ‘W’ 维度。
propertydilation
指定用于扩张卷积的扩张率。
propertygroups
分组卷积将输入分成多少组。输入和输出通道的数量应该能被组数整除。
propertykernel_layout
权重的维度排序。可以是“OIW”、“IOW”等。“O”、“I”、“W”分别代表 num_filter、input_channel 和 width 维度。
propertyout_dtype
输出数据类型,在混合精度设置下设为显式类型。
propertyout_layout
输出维度的排序。可以是‘NCW’、‘NWC’等。'N'、'C'、'W'分别代表批次、通道和宽度维度。默认与输入布局相同。
propertypadding
padding 宽度顺序为(左,右)
Type : 如果填充非零,则输入隐式零填充 Padding 支持对称和非对称作为一整数。
Type : 相同填充应用于两侧两个整数。
propertystrides
指定卷积的步长。
classtvm.relax.op.op_attrs.Conv1DTransposeAttrs
nn.conv1d_transpose 的属性。
propertydata_layout
channel, widthdimensions respectively. Convolution is applied on the ‘W’ dimensions.
输入数据的维度排序。可以是 ‘NCW’、‘NWC’ 等。'N'、'C'、'W' 分别代表批次、通道、宽度。卷积应用于 ‘W’ 维度。
propertydilation
指定用于扩张卷积的扩张率。
propertygroups
分组卷积将输入分成多少组。输入和输出通道的数量应该能被组数整除。
propertykernel_layout
权重的维度排序。可以是“OIW”、“IOW”等。“O”、“I”、“W”分别代表 num_filter、input_channel 和 width 维度。
propertyout_dtype
输出数据类型,在混合精度设置下设为显式类型。
propertyout_layout
输出维度的排序。可以是‘NCW’、‘NWC’等。'N'、'C'、'W'分别代表批次、通道和宽度维度。默认与输入布局相同。
propertyoutput_padding
用于消除输出形状的歧义。
propertypadding
padding 宽度顺序为(左,右)
Type : 如果填充非零,则输入隐式零填充 Padding 支持对称和非对称作为一整数。
Type : 相同填充应用于两侧两个整数。
propertystrides
指定卷积的步长。
classtvm.relax.op.op_attrs.Pool1DAttrs
nn.max_pool1d 和 nn.avg_pool1d 的属性。
propertyceil_mode
一个布尔值,表示是否使用向上取整或向下取整来计算输出形状。使用向上取整时,输入张量的每个元素都将被滑动窗口覆盖。
propertycount_include_pad
当为真时,将包含填充以计算平均值。
propertydilation
指定卷积的膨胀率。
propertylayout
输入数据的维度顺序。可以是“NCW”、“NWC”等。“N”、“C”、“W”分别代表批次、通道和宽度维度。池化操作应用于“W”维度。
propertyout_layout
输出数据的维度顺序。可以是“NCW”、“NWC”等。“N”、“C”、“W”分别代表批次、通道和宽度维度。池化操作应用于“W”维度。
propertypadding
padding 宽度顺序为(左,右)
Type : 如果填充非零,则输入隐式零填充 Padding 支持对称和非对称作为一整数。
Type : 所有边使用相同的填充两个整数。
propertypool_size
池化窗口的大小。
propertystrides
指定卷积的步长。
classtvm.relax.op.op_attrs.Pool3DAttrs
nn.max_pool3d 和 nn.avg_pool3d 的属性。
propertyceil_mode
一个布尔值,表示是否使用向上取整或向下取整来计算输出形状。使用向上取整时,输入张量的每个元素都将被滑动窗口覆盖。
propertycount_include_pad
当为真时,将包含填充以计算平均值。
propertydilation
指定卷积的膨胀率。
propertylayout
输入数据的维度顺序。可以是‘NCDHW’、‘NDHWC’等。‘N’、‘C’、‘D’、‘H’、‘W’分别代表批次、通道、深度、高度和宽度维度。池化操作应用于‘D’、‘H’和‘W’维度。
propertyout_layout
输出数据的维度顺序。可以是‘NCDHW’、‘NDHWC’等。‘N’、‘C’、‘D’、‘H’、‘W’分别代表批次、通道、深度、高度和宽度维度。池化操作应用于‘D’、‘H’和‘W’维度。
propertypadding
back, bottom, right 将使用与 front, top, left 相同的填充 four int : 填充宽度按顺序为(front, top, left, back, bottom, right)。
Type : 如果填充非零,则输入隐式零填充 Padding 支持对称和非对称作为一整数。
Type : 相同填充用于所有侧边三个整数。
propertypool_size
池化窗口的大小。
propertystrides
指定卷积的步长。
classtvm.relax.op.op_attrs.AdaptivePool1DAttrs
1d 自适应池算子的属性。
propertylayout
输入数据的维度顺序。可以是‘NCW’、‘NWC’等。’N’、‘C’、‘W’分别代表批次、通道和宽度维度。池化操作应用于‘W’维度。
propertyout_layout
输出数据的维度顺序。可以是‘NCW’、‘NWC’等。’N’、‘C’、‘W’分别代表批次、通道和宽度维度。池化操作应用于‘W’维度。
propertyoutput_size
输出宽度。
classtvm.relax.op.op_attrs.AdaptivePool3DAttrs
三维自适应池化算子的属性。
propertylayout
输入数据的维度顺序。可以是‘NCDHW’、‘NDHWC’等。‘N’、‘C’、‘D’、‘H’、‘W’分别代表批次、通道、深度、高度和宽度维度。池化操作应用于‘D’、‘H’和‘W’维度。
propertyout_layout
输出数据的维度顺序。可以是‘NCDHW’、‘NDHWC’等。‘N’、‘C’、‘D’、‘H’、‘W’分别代表批次、通道、深度、高度和宽度维度。池化操作应用于‘D’、‘H’和‘W’维度。
propertyoutput_size
输出深度、高度和宽度。
classtvm.relax.op.op_attrs.LeakyReluAttrs
leaky_relu 操作符使用的属性。
propertyalpha
负部分的斜率。
classtvm.relax.op.op_attrs.SoftplusAttrs
softplus 算子中使用的属性。
propertybeta
控制 Softplus 转换锐度的缩放因子。
propertythreshold
确定何时使用线性近似以保证数值稳定性的值。
classtvm.relax.op.op_attrs.PReluAttrs
在 prelu 算子中使用的属性。
propertyaxis
应用 alpha 值的轴。
classtvm.relax.op.op_attrs.PixelShuffleAttrs
pixel_shuffle 算子中使用的属性。
propertyupscale_factor
空间上采样用的缩放因子。
classtvm.relax.op.op_attrs.GroupNormAttrs
group_norm 操作符使用的属性。
propertyaxes
沿其应用归一化(不包括通道轴)的轴。
propertycenter
指示是否将 beta 偏移量添加到归一化张量。
propertychannel_axis
表示通道的轴。
propertyepsilon
为方差添加的小浮点数以避免除以零。
propertynum_groups
将通道分开的组数。
propertyscale
指示是否将 gamma 缩放相乘。
classtvm.relax.op.op_attrs.RMSNormAttrs
rms_norm 操作符使用的属性。
propertyaxes
归一化应用的轴。
propertyepsilon
为方差添加的小浮点数以避免除以零。
classtvm.relax.op.op_attrs.NLLLossAttrs
nll_loss 操作符使用的属性。
propertyignore_index
要忽略的目标值。
propertyreduction
应用于输出的归约方法。可以是'none'、'mean'或'sum'。
classtvm.relax.op.op_attrs.AllReduceAttrs
用于 allreduce 操作符的属性。
propertyin_group
该归约操作是在分组内执行、全局执行还是默认在分组内执行。
propertyop_type
应用于输入数据的归约操作类型。目前仅支持求和。
classtvm.relax.op.op_attrs.AllGatherAttrs
allgather 操作使用的属性。
propertyin_group
allgather 操作是在组内执行还是全局执行,或者默认在组内执行。
propertynum_workers
工作线程的数量,也是给定缓冲区应分块成部分的数量。
classtvm.relax.op.op_attrs.WrapParamAttrs
用于 wrap_param 操作的属性。
propertydtype
目标数据类型。
classtvm.relax.op.op_attrs.QuantizeAttrs
量化/反量化算子中使用的属性。
propertyaxis
通道逐个量化/反量化的输出通道轴。默认值是-1,对应于最后一个轴。
propertyout_dtype
输出数据类型。
classtvm.relax.op.op_attrs.GatherElementsAttrs
操作符的属性。
propertyaxis
用于索引的轴。
classtvm.relax.op.op_attrs.GatherNDAttrs
gather_nd 算子的属性。
propertybatch_dims
批处理维度的数量。
classtvm.relax.op.op_attrs.MeshgridAttrs
网格运算符的属性
propertyindexing
指定网格维度如何排序。
classtvm.relax.op.op_attrs.ScatterElementsAttrs
scatter_elements 操作符的属性。
propertyaxis
选择值的轴。
propertyreduction
scatter_elements 元素的归约模式,可以是“update”、“add”、“mul”、“mean”、“min”或“max”。
classtvm.relax.op.op_attrs.ScatterNDAttrs
scatter_nd 操作符的属性。
propertyreduction
ScatterND 的累积模式,可以是“update”、“add”、“mul”、“min”或“max”。
classtvm.relax.op.op_attrs.SliceScatterAttrs
slice_scatter 操作符的属性。
propertyaxis
插入切片的维度。
classtvm.relax.op.op_attrs.OneHotAttrs
one_hot 操作符的属性。
propertyaxis
Axis to fill. 填充轴。
propertydepth
One-hot 维度的深度。