tvm.topi

TVM 算子库。

TOPI 是 TVM 的算子集合库,为构建计算声明以及优化后的调度提供便利。

某些调度函数可能已针对特定工作负载进行了专门优化。

分析器()

整数算术分析器

Cast(dtype, value[, span])

类型转换表达式。

PrimExpr

所有原始表达式的基类。

函数

abs(x)

逐元素取输入 x 的绝对值。

acos(x)

取输入 x 的反余弦。

acosh(x)

取输入 x 的反双曲余弦。

add(lhs, rhs)

自动广播加法

adv_index(data, indices)

使用张量的 NumPy 风格索引。

all(data[, axis, keepdims])

对给定轴或轴列表的数组元素进行逻辑与运算

any(data[, axis, keepdims])

对给定轴或轴列表的数组元素进行逻辑或运算

arange(start[, stop, step, dtype])

在给定区间内创建具有均匀间隔值的张量。

argmax(data[, axis, keepdims, select_last_index])

返回沿轴的最大值的索引。

argmin(data[, axis, keepdims, select_last_index])

返回沿轴的最小值的索引。

argsort(data[, valid_count, axis, ...])

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

asin(x)

取输入 x 的反正弦。

asinh(x)

取输入 x 的反双曲正弦。

atan(x)

取输入 x 的反正切。

atanh(x)

取输入 x 的反双曲正切。

binary_search(ib, sequence_offset, ...)

用于 CPU 和 GPU 后端的二分搜索的通用 IR 生成器。

bitwise_and(lhs, rhs)

计算数据的逐元素按位与。

bitwise_not(data)

计算数据的逐元素按位非。

bitwise_or(lhs, rhs)

计算数据的逐元素按位或。

bitwise_xor(lhs, rhs)

计算数据的逐元素按位异或。

broadcast_to(data, shape)

将 src 广播到目标形状

cast(x, dtype[, span])

将输入转换为指定的数据类型。

ceil(x)

取输入 x 的 ceiling 值。

ceil_log2(x)

计算整数 ceiling log2,并为 vulkan SPIR-V 提供特殊代码路径,因为它不支持 fp64 上的 log2。

clip(x, a_min, a_max)

裁剪(限制)数组中的值。

collapse_sum(data, target_shape)

返回数据到给定形状的求和。

concatenate(a_tuple[, axis])

沿现有轴连接数组序列。

const_vector(vector[, name])

将常量 numpy 一维向量转换为 tvm 张量

cos(x)

取输入 x 的余弦。

cosh(x)

取输入 x 的双曲余弦。

cumprod(data[, axis, dtype, exclusive])

NumPy 风格的 cumprod 运算。

cumsum(data[, axis, dtype, exclusive])

NumPy 风格的 cumsum 运算。

decl_buffer(shape[, dtype, name, data, ...])

声明一个新的符号缓冲区。

dft(re_data, im_data, inverse)

计算输入的离散傅里叶变换(沿最后一个轴计算)。

div(a, b[, span])

按照 C/C++ 语义计算 a / b。

divide(lhs, rhs)

自动广播除法

dynamic_strided_slice(a, begin, end, ...)

数组切片。

einsum(subscripts, *operand)

对操作数评估爱因斯坦求和约定。

elemwise_sum(xs)

对输入执行逐元素求和

equal(lhs, rhs)

使用自动广播计算 (lhs==rhs)

erf(x)

取输入 x 的高斯误差函数。

exp(x)

取输入 x 的指数。

expand_dims(a, axis[, num_newaxis])

扩展数组的形状。

expand_like(a, shape_like, axis)

使用第二个数组的形状扩展输入数组。

extern(shape, inputs, fcompute[, name, ...])

通过外部函数计算多个张量。

eye(n[, m, k, dtype])

生成单位矩阵或在第 k 条对角线上为 1 的矩阵。

fast_erf(x)

使用 fast_erf 实现取输入 x 的高斯误差函数。

fast_exp(x)

使用 fast_exp 实现取输入 x 的指数

fast_tanh(x)

使用 fast_tanh 实现取输入 x 的双曲正切

fixed_point_multiply(x, multiplier, shift)

数据和定点常数之间的定点乘法,表示为 multiplier * 2^(-shift),其中 multiplier 是一个具有 31 个小数位的 Q 数

fixed_point_multiply_per_axis(x, y, lshift, ...)

数据和定点常数之间的定点乘法,表示为 multiplier * 2^(-shift),其中 multiplier 是一个具有 31 个小数位的 Q 数

flip(a[, axis])

在特定轴中翻转/反转数组元素。

floor(x)

取输入 x 的 floor 值。

floor_divide(lhs, rhs)

自动广播向下除法

floor_mod(lhs, rhs)

自动广播向下取模

floordiv(a, b[, span])

计算两个表达式的向下除法。

floormod(a, b[, span])

计算两个表达式的向下取模。

full(shape, dtype, fill_value)

用 fill_value 填充张量

full_like(x, fill_value)

构造一个与输入张量形状相同的张量,

gather(data, axis, indices)

从给定索引沿给定轴收集值。

gather_nd(a, indices[, batch_dims])

从 n 维数组中收集元素。

get_const_tuple(in_tuple)

验证输入元组是否为 IntImm 或 Var,返回 int 或 Var 元组。

greater(lhs, rhs)

使用自动广播计算 (lhs>rhs)

greater_equal(lhs, rhs)

使用自动广播计算 (lhs>=rhs)

identity(x)

取输入 x 的恒等值。

isfinite(x)

逐元素检查 x 的值是否为有限值。

isinf(x)

逐元素检查 x 的值是否为无穷大。

isnan(x)

逐元素检查 x 的值是否为 NaN。

layout_transform(array, src_layout, dst_layout)

根据 src_layout 和 dst_layout 转换布局

left_shift(lhs, rhs)

自动广播左移

less(lhs, rhs)

使用自动广播计算 (lhs<rhs)

less_equal(lhs, rhs)

使用自动广播计算 (lhs<=rhs)

log(x)

取输入 x 的对数。

log10(x)

取输入 x 的以 10 为底的对数。

log2(x)

取输入 x 的以 2 为底的对数。

logical_and(lhs, rhs)

计算数据的逐元素逻辑与。

logical_not(data)

计算数据的逐元素逻辑非。

logical_or(lhs, rhs)

计算数据的逐元素逻辑或。

logical_xor(lhs, rhs)

计算数据的逐元素逻辑异或。

make_idx(b, e, s, z, i)

返回选择中对应于完整数组中数组位置的数组位置。

matmul(a, b[, transp_a, transp_b])

创建一个运算,该运算计算矩阵乘法(行优先表示法):A(i, k) * B(k, j) 如果 trans_a == trans_b,则为通常的转置组合,否则

matrix_set_diag(data, diagonal[, k, align])

返回一个张量,该张量的输入张量的对角线已替换为提供的对角线值。

max(data[, axis, keepdims])

给定轴或轴列表的数组元素的最大值

maximum(lhs, rhs)

取两个张量的逐元素最大值,带自动广播

meshgrid(a_tuple, indexing)

从坐标向量创建坐标矩阵。

min(data[, axis, keepdims])

给定轴或轴列表的数组元素的最小值

minimum(lhs, rhs)

取两个张量的逐元素最大值,带自动广播

mod(lhs, rhs)

自动广播取模

multiply(lhs, rhs)

自动广播乘法

ndarray_size(array[, dtype])

获取输入数组的元素数量

negative(x)

取输入 x 的负数。

not_equal(lhs, rhs)

使用自动广播计算 (lhs!=rhs)

one_hot(indices, on_value, off_value, depth, ...)

返回一个 one-hot 张量,其中由索引表示的位置取值 on_value,其他位置取值 off_value。

power(lhs, rhs)

自动广播幂运算

prod(data[, axis, keepdims])

给定轴或轴列表的数组元素的乘积

reinterpret(x, dtype)

将输入重新解释为指定的数据类型。

repeat(a, repeats, axis)

重复数组的元素。

reshape(a, newshape)

重塑数组

reverse_sequence(a, seq_lengths[, seq_axis, ...])

反转可变长度切片的张量。

right_shift(lhs, rhs)

自动广播右移

round(x)

将 x 的元素四舍五入到最接近的整数。

rsqrt(x)

取输入 x 的平方根倒数。

scanop(data, binop, identity_value, op_name)

累积二元运算符(扫描),其轴行为类似于 np.cumsum 和 np.cumprod。

scatter_elements(data, indices, updates[, ...])

将更新中的元素分散到复制数据的相应索引中。

scatter_nd(data, indices, updates, mode)

从 n 维数组中分散元素。

searchsorted(sorted_sequence, values[, ...])

查找应插入元素以保持顺序的索引。

sequence_mask(data, valid_length[, ...])

将序列预期长度之外的所有元素设置为常量值。

shape(array[, dtype])

获取输入数组的形状

sigmoid(x)

取输入 x 的 sigmoid tanh。

sign(x)

根据 x 的符号返回 -1、0、1。

sin(x)

取输入 x 的正弦。

sinh(x)

取输入 x 的双曲正弦。

sliding_window(data, axis, window_shape, strides)

在数据张量上滑动窗口。

sort(data[, axis, is_ascend])

沿给定轴执行排序,并返回排序顺序的数组。

sparse_reshape(sparse_indices, prev_shape, ...)

重塑稀疏张量

sparse_to_dense(sparse_indices, ...[, ...])

将稀疏表示形式转换为稠密张量。

split(ary, indices_or_sections[, axis])

将数组拆分为多个子数组。

sqrt(x)

取输入 x 的平方根。

squeeze(a[, axis])

从数组的形状中删除单维度条目。

stack(a, axis)

多次重复整个数组。

stft(data, n_fft, hop_length, win_length, ...)

STFT 计算输入短重叠窗口的傅里叶变换。

strided_set(a, v, begin, end[, strides])

设置数组切片。

strided_slice(a, begin, end[, strides, ...])

数组切片。

subtract(lhs, rhs)

自动广播减法

sum(data[, axis, keepdims])

给定轴或轴列表的数组元素的总和

take(a, indices[, axis, batch_dims, mode])

沿轴从数组中获取元素。

tan(x)

取输入 x 的正切。

tanh(x)

取输入 x 的双曲正切。

tensordot(a, b, axes)

矩阵乘法到张量的推广。

tile(a, reps)

多次重复整个数组。

topk(data[, k, axis, ret_type, is_ascend, dtype])

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

transpose(a[, axes])

置换数组的维度。

trilu(data, k, upper)

给定 2-D 矩阵或批量的 2-D 矩阵,返回张量的上或下三角部分。

trunc(x)

逐元素取输入 x 的截断值。

unravel_index(indices, shape)

将平面索引或平面索引数组转换为坐标数组元组。

where(condition, x, y)

根据条件从 x 或 y 中获取元素。

within_index(b, e, s, i)

返回一个布尔值,指示 i 是否在给定索引内。

异常

InvalidShapeError

topi 函数的无效形状。

class tvm.topi.Analyzer

整数算术分析器

这是一个有状态的分析器类,可用于执行各种符号整数分析。

方法

const_int_bound(expr)

查找 expr 的常量整数边界。

modular_set(expr)

查找 expr 所属的模集合。

simplify(expr[, steps])

通过重写和规范化来简化表达式。

rewrite_simplify(expr)

通过重写规则简化表达式。

canonical_simplify(expr)

通过规范化简化表达式。

int_set(expr, dom_map)

计算涵盖 dom_map 中所有值的 expr 的符号 IntSet。

can_prove(expr[, strength])

检查我们是否可以证明 expr 为真。

bind(var, expr)

将变量绑定到表达式。

constraint_scope(constraint)

创建约束范围。

update(var, info[, override])

更新有关 var 的信息

can_prove_equal(lhs, rhs)

我们是否可以证明 lhs == rhs

属性

enabled_extensions

返回当前启用的扩展

const_int_bound(expr)

查找 expr 的常量整数边界。

参数:

expr (PrimExpr) – 表达式。

返回:

bound – 结果边界

返回类型:

ConstIntBound

modular_set(expr)

查找 expr 所属的模集合。

参数:

expr (PrimExpr) – 表达式。

返回:

result – 结果。

返回类型:

ModularSet

simplify(expr, steps=2)

通过重写和规范化来简化表达式。

参数:
  • expr (PrimExpr) – 表达式。

  • steps (简化按以下顺序运行) – rewrite_simplify(步骤 1)-> canonical_simplify(步骤 2)-> rewrite_simplify(步骤 3)-> canonical_simplify(步骤 4)-> … 参数 steps 控制运行多少步骤。 默认为 2,即 rewrite_simplify + canonical_simplify。

返回:

result – 结果。

返回类型:

Expr

rewrite_simplify(expr)

通过重写规则简化表达式。

参数:

expr (PrimExpr) – 表达式。

返回:

result – 结果。

返回类型:

Expr

canonical_simplify(expr)

通过规范化简化表达式。

参数:

expr (PrimExpr) – 表达式。

返回:

result – 结果。

返回类型:

Expr

int_set(expr, dom_map)

计算涵盖 dom_map 中所有值的 expr 的符号 IntSet。

参数:
  • expr (PrimExpr) – 表达式。

  • dom_map (Dict[tvm.tir.Var, tvm.arith.IntSet]) – 要放松的变量的域。

返回:

result – 结果。

返回类型:

IntSet

can_prove(expr, strength=ProofStrength.DEFAULT)

检查我们是否可以证明 expr 为真。

参数:
  • expr (PrimExpr) – 表达式。

  • strength (ProofStrength) – 证明强度

返回:

result – 结果。

返回类型:

Expr

bind(var: Var, expr: PrimExpr | Range)

将变量绑定到表达式。

参数:
  • var (tvm.tir.Var) – 变量。

  • expr (Union[tir.PrimExpr, ir.Range]) – 要绑定到的表达式或范围。

constraint_scope(constraint)

创建约束范围。

参数:

constraint (PrimExpr) – 约束表达式。

返回:

scope – 约束作用域

返回类型:

ConstraintScope

示例

x = te.var("x")
analyzer = tvm.arith.Analyzer()
with analzyer.constraint_scope(x % 3 == 0):
    # constraint in effect
    assert analyzer.modular_set(x).coeff == 3
# constraint no longer in effect
assert analyzer.modular_set(x).coeff != 3
update(var, info, override=False)

更新有关 var 的信息

参数:
  • var (tvm.tir.Var) – 变量。

  • info (tvm.Object) – 相关信息。

  • override (bool) – 是否允许覆盖。

can_prove_equal(lhs: PrimExpr, rhs: PrimExpr)

我们是否可以证明 lhs == rhs

参数:
返回:

result – 我们是否可以证明 lhs == rhs

返回类型:

bool

property enabled_extensions: Extension

返回当前启用的扩展

class tvm.topi.Cast(dtype, value, span: Span | None = None)

类型转换表达式。

参数:
  • dtype (str) – 数据类型。

  • value (PrimExpr) – 函数的值。

  • span (Optional[Span]) – 此表达式在源代码中的位置。

class tvm.topi.PrimExpr

所有原始表达式的基类。

PrimExpr 用于底层代码优化和整数分析。

tvm.topi.abs(x)

逐元素取输入 x 的绝对值。

参数:

x (tvm.te.Tensor) – 输入参数。

返回:

y – 结果。

返回类型:

tvm.te.Tensor

tvm.topi.acos(x)

取输入 x 的反余弦。

参数:

x (tvm.te.Tensor) – 输入参数。

返回:

y – 结果。

返回类型:

tvm.te.Tensor

tvm.topi.acosh(x)

取输入 x 的反双曲余弦。

参数:

x (tvm.te.Tensor) – 输入参数。

返回:

y – 结果。

返回类型:

tvm.te.Tensor

tvm.topi.add(lhs, rhs)

自动广播加法

参数:
返回:

ret – 如果两个操作数都是 Expr,则返回 Expr。否则返回 Tensor。

返回类型:

tvm.te.Tensor 或 Expr

tvm.topi.adv_index(data, indices)

使用张量的 NumPy 风格索引。

参数:
返回:

result – 输出张量

返回类型:

tvm.te.Tensor

tvm.topi.all(data, axis=None, keepdims=False)

对给定轴或轴列表的数组元素进行逻辑与运算

参数:
  • data (tvm.te.Tensor) – 输入 tvm 布尔张量

  • axis (Noneinttuple of int) – 执行逻辑与运算的轴或轴组。默认值 axis=None 将对输入数组的所有元素执行逻辑与运算。如果 axis 为负数,则从最后一个轴向第一个轴计数。

  • keepdims (bool) – 如果设置为 True,则被缩减的轴将作为大小为 1 的维度保留在结果中。使用此选项,结果将与输入数组正确广播。

返回:

ret

返回类型:

tvm.te.Tensor

tvm.topi.any(data, axis=None, keepdims=False)

对给定轴或轴列表的数组元素进行逻辑或运算

参数:
  • data (tvm.te.Tensor) – 输入 tvm 布尔张量

  • axis (Noneinttuple of int) – 执行逻辑或运算的轴或轴组。默认值 axis=None 将对输入数组的所有元素执行逻辑或运算。如果 axis 为负数,则从最后一个轴向第一个轴计数。

  • keepdims (bool) – 如果设置为 True,则被缩减的轴将作为大小为 1 的维度保留在结果中。使用此选项,结果将与输入数组正确广播。

返回:

ret

返回类型:

tvm.te.Tensor

tvm.topi.arange(start, stop=None, step=1, dtype='float32')

在给定区间内创建具有均匀间隔值的张量。

参数:
  • start (tvm.Expr, optional) – 区间的起始值。区间包含此值。默认起始值为 0。

  • stop (tvm.Expr) – 区间的终止值。区间不包含此值。

  • step (tvm.Expr, optional) – 值之间的间距。默认步长为 1。

  • dtype (str, optional) – 目标数据类型。

返回:

result – 结果张量。

返回类型:

tvm.te.Tensor

tvm.topi.argmax(data, axis=None, keepdims=False, select_last_index=False)

返回沿轴的最大值的索引。

参数:
  • data (tvm.te.Tensor) – 输入 tvm 张量

  • axis (Noneinttuple of int) – 执行 argmax 操作的轴或轴组。默认值 axis=None 将查找输入数组元素中最大元素的索引。如果 axis 为负数,则从最后一个轴向第一个轴计数。

  • keepdims (bool) – 如果设置为 True,则被缩减的轴将作为大小为 1 的维度保留在结果中。使用此选项,结果将与输入数组正确广播。

  • select_last_index (bool) – 如果最大元素多次出现,是否选择最后一个索引,否则选择第一个索引。

返回:

ret

返回类型:

tvm.te.Tensor

tvm.topi.argmin(data, axis=None, keepdims=False, select_last_index=False)

返回沿轴的最小值的索引。

参数:
  • data (tvm.te.Tensor) – 输入 tvm 张量

  • axis (Noneinttuple of int) – 执行 argmin 操作的轴或轴组。默认值 axis=None 将查找输入数组所有元素中最小元素的索引。如果 axis 为负数,则从最后一个轴向第一个轴计数。

  • keepdims (bool) – 如果设置为 True,则被缩减的轴将作为大小为 1 的维度保留在结果中。使用此选项,结果将与输入数组正确广播。

  • select_last_index (bool) – 如果最小元素多次出现,是否选择最后一个索引,否则选择第一个索引。

返回:

ret

返回类型:

tvm.te.Tensor

tvm.topi.argsort(data, valid_count=None, axis=-1, is_ascend=1, dtype='float32')

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

参数:
  • data (tvm.te.Tensor) – 输入张量。

  • valid_count (tvm.te.Tensor, optional) – 1-D 张量,表示有效框的数量。

  • axis (int, optional) – 对输入张量进行排序的轴。默认情况下使用展平的数组。

  • is_ascend (boolean, optional) – 是否按升序或降序排序。

  • dtype (string, optional) – 输出索引的数据类型。

返回:

out – 排序后的索引张量。

返回类型:

tvm.te.Tensor

示例

# An example to use argsort
dshape = (1, 5, 6)
data = te.placeholder(dshape, name="data")
axis = 0
is_ascend = False
out = argsort(data, axis=axis, is_ascend=is_ascend)
np_data = np.random.uniform(dshape)
s = topi.generic.schedule_argsort(out)
f = tvm.compile(s, [data, out], "llvm")
dev = tvm.cpu()
tvm_data = tvm.nd.array(np_data, dev)
tvm_out = tvm.nd.array(np.zeros(dshape, dtype=data.dtype), dev)
f(tvm_data, tvm_out)
tvm.topi.asin(x)

取输入 x 的反正弦。

参数:

x (tvm.te.Tensor) – 输入参数。

返回:

y – 结果。

返回类型:

tvm.te.Tensor

tvm.topi.asinh(x)

取输入 x 的反双曲正弦。

参数:

x (tvm.te.Tensor) – 输入参数。

返回:

y – 结果。

返回类型:

tvm.te.Tensor

tvm.topi.atan(x)

取输入 x 的反正切。

参数:

x (tvm.te.Tensor) – 输入参数。

返回:

y – 结果。

返回类型:

tvm.te.Tensor

tvm.topi.atanh(x)

取输入 x 的反双曲正切。

参数:

x (tvm.te.Tensor) – 输入参数。

返回:

y – 结果。

返回类型:

tvm.te.Tensor

tvm.topi.binary_search(ib, sequence_offset, search_range, sorted_sequence, value, right, out_dtype)

用于 CPU 和 GPU 后端的二分搜索的通用 IR 生成器。

sorted_sequence 是一个 N 维缓冲区,我们希望在其最内层维度中搜索 value,而 search_range 是最内层维度的大小。sequence_offset 是一个 1 维线性化偏移量,指定要搜索的最内层序列。

因此,对 value 的搜索将在 sorted_sequence[sequence_offset:(sequence_offset + search_range)] 上执行。请注意,我们通过 1 维线性化索引来索引 N 维缓冲区。

tvm.topi.bitwise_and(lhs, rhs)

计算数据的逐元素按位与。

参数:
返回:

ret – 如果两个操作数都是 Expr,则返回 Expr。否则返回 Tensor。

返回类型:

tvm.te.Tensor 或 Expr

tvm.topi.bitwise_not(data)

计算数据的逐元素按位非。

参数:

data (tvm.te.TensorExpr)

返回:

ret – 如果操作数是 Expr,则返回 Expr。否则返回 Tensor。

返回类型:

tvm.te.Tensor 或 Expr

tvm.topi.bitwise_or(lhs, rhs)

计算数据的逐元素按位或。

参数:
返回:

ret – 如果两个操作数都是 Expr,则返回 Expr。否则返回 Tensor。

返回类型:

tvm.te.Tensor 或 Expr

tvm.topi.bitwise_xor(lhs, rhs)

计算数据的逐元素按位异或。

参数:
返回:

ret – 如果两个操作数都是 Expr,则返回 Expr。否则返回 Tensor。

返回类型:

tvm.te.Tensor 或 Expr

tvm.topi.broadcast_to(data, shape)

将 src 广播到目标形状

我们遵循 numpy 广播规则。另请参阅 https://docs.scipy.org.cn/doc/numpy/user/basics.broadcasting.html

参数:
返回:

ret

返回类型:

tvm.te.Tensor

tvm.topi.cast(x, dtype, span=None)

将输入转换为指定的数据类型。

参数:
  • x (tvm.te.TensorExpr) – 输入参数。

  • dtype (str) – 数据类型。

  • span (Optional[Span]) – 强制类型转换在源代码中的位置。

返回:

y – 结果。

返回类型:

tvm.te.Tensor

tvm.topi.ceil(x)

取输入 x 的 ceiling 值。

参数:

x (tvm.te.Tensor) – 输入参数。

返回:

y – 结果。

返回类型:

tvm.te.Tensor

tvm.topi.ceil_log2(x)

计算整数 ceil log2,对于 vulkan SPIR-V 有特殊的代码路径,因为它不支持 fp64 上的 log2。相反,当目标是 vulkan 时,我们通过 clz intrinsic 计算整数 ceil_log2。

参数:

x (tvm.te.Tensor) – 输入参数。

返回:

y – 结果。

返回类型:

tvm.te.Tensor

tvm.topi.clip(x, a_min, a_max)

裁剪(限制)数组中的值。给定一个区间,区间外的值将被裁剪到区间边缘。

参数:
  • x (tvm.te.Tensor) – 输入参数。

  • a_min (tvm.tir.PrimExpr) – 最小值。

  • a_max (tvm.tir.PrimExpr) – 最大值。

返回:

y – 结果。

返回类型:

tvm.te.Tensor

tvm.topi.collapse_sum(data, target_shape)

返回数据到给定形状的求和。

collapse_sum 旨在作为自动微分过程中 topi 广播运算符的反向运算符。

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

在计算期间,从右到左检查 data.shape 和 target_shape 的轴。对于每个轴,如果它满足以下任一条件:- 存在于 data 中但不存在于 target_shape 中,或者 - 在 data 中大于 1 且在 target_shape 中等于 1,则 data 将沿此轴求和。

参数:
返回:

ret – 求和后的结果张量。

返回类型:

tvm.te.Tensor

tvm.topi.concatenate(a_tuple, axis=0)

沿现有轴连接数组序列。

参数:
  • a_tuple (tuple of tvm.te.Tensor) – 要连接的数组

  • axis (int, optional) – 数组将沿其连接的轴。默认为 0。

返回:

ret

返回类型:

tvm.te.Tensor

tvm.topi.const_vector(vector, name='const_vector')

将常量 numpy 一维向量转换为 tvm 张量

参数:
  • vector (numpy.ndarray) – 常量输入数组

  • name (str, optional) – 输出算子的名称

返回:

tensor – 创建的张量

返回类型:

Tensor

tvm.topi.cos(x)

取输入 x 的余弦。

参数:

x (tvm.te.Tensor) – 输入参数。

返回:

y – 结果。

返回类型:

tvm.te.Tensor

tvm.topi.cosh(x)

取输入 x 的双曲余弦。

参数:

x (tvm.te.Tensor) – 输入参数。

返回:

y – 结果。

返回类型:

tvm.te.Tensor

tvm.topi.cumprod(data: Tensor, axis: int | None = None, dtype: int | None = None, exclusive: bool | None = None) Tensor

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

参数:
  • data (tvm.te.Tensor) – 运算符的输入数据。

  • axis (int, optional) – 计算累积乘积的轴。默认值 (None) 是计算展平数组上的累积乘积。

  • dtype (string, optional) – 返回数组以及在其中将元素相乘的累加器类型。如果未指定 dtype,则默认为 data 的 dtype。

  • exclusive (bool, optional) – 如果为 True,将返回独占乘积,其中不包括第一个元素。换句话说,如果为 True,则第 j 个输出元素将是前 (j-1) 个元素的乘积。否则,它将是前 j 个元素的乘积。

返回:

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

返回类型:

tvm.te.Tensor

tvm.topi.cumsum(data: Tensor, axis: int | None = None, dtype: str | None = None, exclusive: bool | None = None) Tensor

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

参数:
  • data (tvm.te.Tensor) – 运算符的输入数据。

  • axis (int, optional) – 计算累积和的轴。默认值 (None) 是计算展平数组上的累积和。

  • dtype (string, optional) – 返回数组以及在其中将元素相加的累加器类型。如果未指定 dtype,则默认为 data 的 dtype。

  • exclusive (bool, optional) – 如果为 True,将返回独占和,其中不包括第一个元素。换句话说,如果为 True,则第 j 个输出元素将是前 (j-1) 个元素的和。否则,它将是前 j 个元素的和。

返回:

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

返回类型:

tvm.te.Tensor

tvm.topi.decl_buffer(shape, dtype=None, name='buffer', data=None, strides=None, elem_offset=None, scope='', data_alignment=-1, offset_factor=0, buffer_type='', axis_separators=None, span=None)

声明一个新的符号缓冲区。

通常情况下,缓冲区在 lower 和 build 期间自动创建。仅当用户想要指定自己的缓冲区布局时才需要这样做。

有关缓冲区用法的详细讨论,请参阅下面的注释。

参数:
  • shape (tuple of Expr) – 缓冲区的形状。

  • dtype (str, optional) – 缓冲区的数据类型。

  • name (str, optional) – 缓冲区的名称。

  • data (tir.Var, optional) – 缓冲区中的数据指针。

  • strides (array of Expr) – 缓冲区的步幅。

  • elem_offset (Expr, 可选) – 数组数据起始偏移量。以 dtype 的元素数量为单位。

  • scope (str, 可选) – 缓冲区的存储范围,如果不是全局的。如果 scope 等于空字符串,则表示它是全局内存。

  • data_alignment (int, 可选) – 数据指针的对齐字节数。如果传递 -1,则对齐方式将设置为 TVM 的内部默认值。

  • offset_factor (int, 可选) – elem_offset 字段的因子,设置后,elem_offset 必须是 offset_factor 的倍数。如果传递 0,则对齐方式将设置为 1。如果传递非零值,如果 elem_offset 不为 None,我们将为 elem_offset 创建一个 Var。

  • buffer_type (str, 可选, {"", "auto_broadcast"}) – auto_broadcast 缓冲区允许实现广播计算,而无需考虑维度大小是否等于 1。如果维度 j 的形状等于 1,TVM 会将 buffer[i][j][k] 映射到 buffer[i][0][k]。

  • axis_separators (list of int, 可选) – 如果传递,则为轴组之间分隔符的列表,每个轴组将展平为一个输出轴。对于平面内存空间,应为 None 或空列表。

  • span (Optional[Span]) – decl_buffer 创建在源代码中的位置。

返回:

buffer – 创建的缓冲区

返回类型:

tvm.tir.Buffer

注意

Buffer 数据结构反映了 dlpack 中的 DLTensor 结构。虽然 DLTensor 数据结构非常通用,但通常创建仅处理数据结构特定情况的函数并使编译后的函数从中受益是很有帮助的。

如果用户在构造函数时传递了 strides 并且 elem_offset 为 None,则该函数将专门用于紧凑且对齐的 DLTensor。如果用户将完全通用的符号数组传递给 strides,则生成的函数将变为完全通用的。

tvm.topi.dft(re_data: Tensor, im_data: Tensor, inverse: IntImm)

计算输入的离散傅里叶变换(沿最后一个轴计算)。这给出了信号的频率分量随时间的变化。

参数:
  • re_data (te.Tensor) – N 维张量,输入信号的实部。

  • im_data (te.Tensor) – N 维张量,输入信号的虚部。如果信号是实数,则此张量的值为零。

  • inverse (bool) – 是否执行逆离散傅里叶变换。

返回:

  • re_output (te.Tensor) – 输入的傅里叶变换(实部)。

  • im_output (te.Tensor) – 输入的傅里叶变换(虚部)。

tvm.topi.div(a, b, span=None)

按照 C/C++ 语义计算 a / b。

参数:
  • a (PrimExpr) – 左侧操作数,已知为非负数。

  • b (PrimExpr) – 右侧操作数,已知为非负数。

  • span (Optional[Span]) – 此运算符在源代码中的位置。

返回:

res – 结果表达式。

返回类型:

PrimExpr

注意

当操作数是整数时,返回 truncdiv(a, b, span)。

tvm.topi.divide(lhs, rhs)

自动广播除法

参数:
返回:

ret – 如果两个操作数都是 Expr,则返回 Expr。否则返回 Tensor。

返回类型:

tvm.te.Tensor 或 Expr

tvm.topi.dynamic_strided_slice(a, begin, end, strides, output_shape)

数组切片。

参数:
  • a (tvm.te.Tensor) – 要切片的张量。

  • begin (tvm.te.Tensor) – 切片起始索引。

  • end (tvm.te.Tensor) – 指示切片结束的索引。

  • strides (tvm.te.Tensor) – 指定步长值,它可以为负数,在这种情况下,输入张量将在该特定轴上反转。

  • output_shape (list of PrimExpr) – 指定输出形状

返回:

ret

返回类型:

tvm.te.Tensor

tvm.topi.einsum(subscripts, *operand)

对操作数评估爱因斯坦求和约定。

参数:
  • subscripts (string) – 指定用于求和的下标,以逗号分隔的下标标签列表形式。除非包含显式指示符 ‘->’ 以及精确输出形式的下标标签,否则将执行隐式(经典爱因斯坦求和)计算。

  • a_tuple (tuple of tvm.te.Tensor) – 这些是运算的张量。tvm 中 einsum 与 numpy 的唯一区别是它需要张量的额外括号。例如,topi.einsum(“ij, jk -> ik”, (A, B))。

返回:

out – 基于爱因斯坦求和约定的计算结果。

返回类型:

tvm.te.Tensor

tvm.topi.elemwise_sum(xs)

对输入执行逐元素求和

参数:

xs (list of tvm.te.Tensor) – 输入参数。

返回:

y – 结果。

返回类型:

tvm.te.Tensor

tvm.topi.equal(lhs, rhs)

使用自动广播计算 (lhs==rhs)

参数:
返回:

ret – 如果两个操作数都是 Expr,则返回 Expr。否则返回 Tensor。

返回类型:

tvm.te.Tensor 或 Expr

tvm.topi.erf(x)

取输入 x 的高斯误差函数。

参数:

x (tvm.te.Tensor) – 输入参数。

返回:

y – 结果。

返回类型:

tvm.te.Tensor

tvm.topi.exp(x)

取输入 x 的指数。

参数:

x (tvm.te.Tensor) – 输入参数。

返回:

y – 结果。

返回类型:

tvm.te.Tensor

tvm.topi.expand_dims(a, axis, num_newaxis=1)

扩展数组的形状。

参数:
  • a (tvm.te.Tensor) – 要扩展的张量。

  • num_newaxis (int, 可选) – 要在轴上插入的 newaxis 的数量

返回:

ret

返回类型:

tvm.te.Tensor

tvm.topi.expand_like(a, shape_like, axis)

使用第二个数组的形状扩展输入数组。此操作始终可以由在那些未压缩轴上进行 unsqueezing 和 expanding dims 组成。

示例

input = [ 12.  19.  27.]
input.shape = (3,)

new_shape_array = [[[1,2],[2,3],[1,3]],
                [[1,4],[4,3],[5,2]],
                [[7,1],[7,2],[7,3]]]
new_shape_array.shape = (3, 3, 2)

expand_like(input, [1,2], new_shape_array) =
                [[[12,12],[12,12],[12,12]],
                [[19,19],[19,19],[19,19]],
                [[27,27],[27,27],[27,27]]]
参数:
返回:

ret

返回类型:

tvm.te.Tensor

tvm.topi.extern(shape, inputs, fcompute, name='extern', dtype=None, in_buffers=None, out_buffers=None, tag='', attrs=None)

通过外部函数计算多个张量。

参数:
  • shape (tuple or list of tuples.) – 输出的形状。

  • inputs (list of Tensor) – 输入

  • fcompute (inputs, outputs-> stmt 的 lambda 函数) –

    指定执行计算的 IR 语句。有关 fcompute 的函数签名,请参见以下注释

    注意

    参数

    返回

  • name (str, 可选) – 张量的名称提示

  • dtype (str or list of str, 可选) – 输出的数据类型,默认情况下 dtype 将与输入相同。

  • in_buffers (tvm.tir.Buffer or list of tvm.tir.Buffer, 可选) – 输入缓冲区。

  • out_buffers (tvm.tir.Buffer or list of tvm.tir.Buffer, 可选) – 输出缓冲区。

tag: str, optional

有关计算的其他标签信息。

attrs: dict, optional

有关计算的其他辅助属性。

返回:

tensor – 创建的张量或包含多个输出的张量元组。

返回类型:

Tensorlist of Tensors

示例

在下面的代码中,C 是通过调用外部 PackedFunc `tvm.contrib.cblas.matmul` 生成的

A = te.placeholder((n, l), name="A")
B = te.placeholder((l, m), name="B")
C = te.extern((n, m), [A, B],
               lambda ins, outs: tvm.tir.call_packed(
                  "tvm.contrib.cblas.matmul",
                    ins[0], ins[1], outs[0], 0, 0), name="C")
tvm.topi.eye(n: int, m: int | None = None, k: int = 0, dtype: str = 'float32') Tensor

生成单位矩阵或在第 k 条对角线上为 1 的矩阵。

参数:
  • n (int) – 行数

  • m (int, 可选) – 列数。如果为 None,则默认为 n。

  • k (int, 可选) – 对角线的索引。0(默认值)表示主对角线。正值表示上对角线,负值表示下对角线。

  • dtype (str, 可选) – 返回数组的数据类型。

返回:

y – 结果。

返回类型:

tvm.te.Tensor

tvm.topi.fast_erf(x)

使用 fast_erf 实现取输入 x 的高斯误差函数。

参数:

x (tvm.te.Tensor) – 输入参数。

返回:

y – 结果。

返回类型:

tvm.te.Tensor

tvm.topi.fast_exp(x)

使用 fast_exp 实现取输入 x 的指数

参数:

x (tvm.te.Tensor) – 输入参数。

返回:

y – 结果。

返回类型:

tvm.te.Tensor

tvm.topi.fast_tanh(x)

使用 fast_tanh 实现取输入 x 的双曲正切

参数:

x (tvm.te.Tensor) – 输入参数。

返回:

y – 结果。

返回类型:

tvm.te.Tensor

tvm.topi.fixed_point_multiply(x, multiplier, shift)

数据和定点常数之间的定点乘法,表示为 multiplier * 2^(-shift),其中 multiplier 是一个具有 31 个小数位的 Q 数

参数:
  • x (tvm.te.TensorExpr) – 输入参数。

  • multiplier (int) – 固定浮点数的乘数,描述为 multiplier*2^(-shift)。

  • shift (int) – 固定浮点数的移位,描述为 multiplier*2^(-shift)。

返回:

y – 结果。

返回类型:

tvm.te.Tensor

tvm.topi.fixed_point_multiply_per_axis(x: Tensor, y: Tensor, lshift: Tensor, rshift: Tensor, is_lshift_required: int, is_rshift_required: int, axes)

数据和定点常数之间的定点乘法,表示为 multiplier * 2^(-shift),其中 multiplier 是一个具有 31 个小数位的 Q 数

参数:
  • x (tvm.te.Tensor) – 输入参数。

  • y (tvm.te.Tensor) – 固定浮点数的乘数,描述为 multiplier*2^(-shift)。

  • lshift (tvm.te.Tensor) – 固定浮点数的左移位数,描述为 multiplier*2^(-shift)。

  • rshift (tvm.te.Tensor) – 固定浮点数的右移位数,描述为 multiplier*2^(-shift)。

  • is_lshift_required (int) – 是否需要进行左移。

  • is_rshift_required (int) – 是否需要进行右移。

返回:

z – 结果。

返回类型:

tvm.te.Tensor

tvm.topi.flip(a, axis=0)

在特定轴中翻转/反转数组元素。

参数:
  • a (tvm.te.Tensor) – 要扩展的张量。

  • axis (int, 可选) – 张量将沿其反转的轴。

返回:

ret

返回类型:

tvm.te.Tensor

tvm.topi.floor(x)

取输入 x 的 floor 值。

参数:

x (tvm.te.Tensor) – 输入参数。

返回:

y – 结果。

返回类型:

tvm.te.Tensor

tvm.topi.floor_divide(lhs, rhs)

自动广播向下除法

参数:
返回:

ret – 如果两个操作数都是 Expr,则返回 Expr。否则返回 Tensor。

返回类型:

tvm.te.Tensor 或 Expr

tvm.topi.floor_mod(lhs, rhs)

自动广播向下取模

参数:
返回:

ret – 如果两个操作数都是 Expr,则返回 Expr。否则返回 Tensor。

返回类型:

tvm.te.Tensor 或 Expr

tvm.topi.floordiv(a, b, span=None)

计算两个表达式的向下除法。

参数:
  • a (PrimExpr) – 左侧操作数

  • b (PrimExpr) – 右侧操作数

  • span (Optional[Span]) – 此运算符在源代码中的位置。

返回:

res – 结果表达式。

返回类型:

PrimExpr

tvm.topi.floormod(a, b, span=None)

计算两个表达式的向下取模。

参数:
  • a (PrimExpr) – 左侧操作数

  • b (PrimExpr) – 右侧操作数

  • span (Optional[Span]) – 此运算符在源代码中的位置。

返回:

res – 结果表达式。

返回类型:

PrimExpr

tvm.topi.full(shape, dtype, fill_value)

用 fill_value 填充张量

参数:
  • shape (tuple) – 输入张量形状。

  • dtype (str) – 数据类型

  • fill_value (float) – 要填充的值

返回:

y – 结果。

返回类型:

tvm.te.Tensor

tvm.topi.full_like(x, fill_value)
构造一个与输入张量形状相同的张量,

然后用 fill_value 填充张量。

参数:
返回:

y – 结果。

返回类型:

tvm.te.Tensor

tvm.topi.gather(data, axis, indices)

从给定索引沿给定轴收集值。

例如,对于 3D 张量,输出计算如下

out[i][j][k] = data[indices[i][j][k]][j][k]  # if axis == 0
out[i][j][k] = data[i][indices[i][j][k]][k]  # if axis == 1
out[i][j][k] = data[i][j][indices[i][j][k]]  # if axis == 2

indices 必须与 data 具有相同的形状,除了维度 axis 必须是非空的。输出将与 indices 具有相同的形状。

参数:
  • data (tvm.te.Tensor) – 运算符的输入数据。

  • axis (int) – 索引所沿的轴。

  • indices (tvm.te.Tensor) – 要提取的值的索引。

返回:

ret

返回类型:

tvm.te.Tensor

tvm.topi.gather_nd(a, indices, batch_dims=0)

从 n 维数组中收集元素。

参数:
返回:

ret

返回类型:

tvm.te.Tensor

tvm.topi.get_const_tuple(in_tuple)

验证输入元组是否为 IntImm 或 Var,返回 int 或 Var 元组。

参数:

in_tuple (tuple of Expr) – 输入。

返回:

out_tuple – 输出。

返回类型:

tuple of int

tvm.topi.greater(lhs, rhs)

使用自动广播计算 (lhs>rhs)

参数:
返回:

ret – 如果两个操作数都是 Expr,则返回 Expr。否则返回 Tensor。

返回类型:

tvm.te.Tensor 或 Expr

tvm.topi.greater_equal(lhs, rhs)

使用自动广播计算 (lhs>=rhs)

参数:
返回:

ret – 如果两个操作数都是 Expr,则返回 Expr。否则返回 Tensor。

返回类型:

tvm.te.Tensor 或 Expr

tvm.topi.identity(x)

取输入 x 的恒等值。

参数:

x (tvm.te.Tensor) – 输入参数。

返回:

y – 结果。

返回类型:

tvm.te.Tensor

tvm.topi.isfinite(x)

逐元素检查 x 的值是否为有限值。

参数:

x (tvm.te.Tensor) – 输入参数。

返回:

y – 结果。

返回类型:

tvm.te.Tensor

tvm.topi.isinf(x)

逐元素检查 x 的值是否为无穷大。

参数:

x (tvm.te.Tensor) – 输入参数。

返回:

y – 结果。

返回类型:

tvm.te.Tensor

tvm.topi.isnan(x)

逐元素检查 x 的值是否为 NaN。

参数:

x (tvm.te.Tensor) – 输入参数。

返回:

y – 结果。

返回类型:

tvm.te.Tensor

tvm.topi.layout_transform(array, src_layout, dst_layout, schedule_rule='None')

根据 src_layout 和 dst_layout 转换布局

参数:
  • array (tvm.te.Tensor) – 源数组。

  • src_layout (str) – 源布局。

  • dst_layout (str) – 目标布局。

  • schedule_rule (str) – 要应用的调度规则(如果有)

tvm.topi.left_shift(lhs, rhs)

自动广播左移

参数:
返回:

ret – 如果两个操作数都是 Expr,则返回 Expr。否则返回 Tensor。

返回类型:

tvm.te.Tensor 或 Expr

tvm.topi.less(lhs, rhs)

使用自动广播计算 (lhs<rhs)

参数:
返回:

ret – 如果两个操作数都是 Expr,则返回 Expr。否则返回 Tensor。

返回类型:

tvm.te.Tensor 或 Expr

tvm.topi.less_equal(lhs, rhs)

使用自动广播计算 (lhs<=rhs)

参数:
返回:

ret – 如果两个操作数都是 Expr,则返回 Expr。否则返回 Tensor。

返回类型:

tvm.te.Tensor 或 Expr

tvm.topi.log(x)

取输入 x 的对数。

参数:

x (tvm.te.Tensor) – 输入参数。

返回:

y – 结果。

返回类型:

tvm.te.Tensor

tvm.topi.log10(x)

取输入 x 的以 10 为底的对数。

参数:

x (tvm.te.Tensor) – 输入参数。

返回:

y – 结果。

返回类型:

tvm.te.Tensor

tvm.topi.log2(x)

取输入 x 的以 2 为底的对数。

参数:

x (tvm.te.Tensor) – 输入参数。

返回:

y – 结果。

返回类型:

tvm.te.Tensor

tvm.topi.logical_and(lhs, rhs)

计算数据的逐元素逻辑与。

参数:
返回:

ret – 如果两个操作数都是 Expr,则返回 Expr。否则返回 Tensor。

返回类型:

tvm.te.Tensor 或 Expr

tvm.topi.logical_not(data)

计算数据的逐元素逻辑非。

参数:

data (tvm.te.TensorExpr)

返回:

ret – 如果操作数是 Expr,则返回 Expr。否则返回 Tensor。

返回类型:

tvm.te.Tensor 或 Expr

tvm.topi.logical_or(lhs, rhs)

计算数据的逐元素逻辑或。

参数:
返回:

ret – 如果两个操作数都是 Expr,则返回 Expr。否则返回 Tensor。

返回类型:

tvm.te.Tensor 或 Expr

tvm.topi.logical_xor(lhs, rhs)

计算数据的逐元素逻辑异或。

参数:
返回:

ret – 如果两个操作数都是 Expr,则返回 Expr。否则返回 Tensor。

返回类型:

tvm.te.Tensor 或 Expr

tvm.topi.make_idx(b, e, s, z, i)

返回选择中对应于完整数组中数组位置的数组位置。

返回的值仅在对于相同的参数集 within_index() 返回 True 时才有意义。

参数:
  • b (Expr) – 索引的起始位置

  • e (Expr) – 索引的结束位置

  • s (Expr) – 索引的步幅

  • z (Expr) – 被索引维度的尺寸

  • i (Expr) – 数组位置

返回:

position – 对应于选择中数组位置的整数表达式。

返回类型:

Expr

tvm.topi.matmul(a, b, transp_a=False, transp_b=False)

创建一个运算,该运算计算矩阵乘法(行优先表示法):A(i, k) * B(k, j) 如果 trans_a == trans_b,则为通常的转置组合,否则

参数:
  • a (矩阵 A)

  • b (矩阵 B)

  • trans_a (A 的布局是否转置?)

  • trans_b (B 的布局是否转置?)

返回类型:

一个 Tensor,其 op 成员是 matmul 操作

tvm.topi.matrix_set_diag(data, diagonal, k=0, align='RIGHT_LEFT')

返回一个张量,该张量的输入张量的对角线已替换为提供的对角线值。

参数:
  • data (tvm.te.Tensor) – 输入 Tensor。

  • diagonal (tvm.te.Tensor) – 要填充到对角线的值。

  • k (inttuple 类型的 int,可选) – 对角线偏移量。要设置的对角线或对角线范围。(默认为 0)正值表示超对角线,0 表示主对角线,负值表示次对角线。k 可以是单个整数(用于单个对角线)或一对整数,指定矩阵带的下限和上限。k[0] 不得大于 k[1]。

  • align (string,可选) – 某些对角线比 max_diag_len 短,需要填充。align 是一个字符串,指定应如何分别对齐超对角线和次对角线。有四种可能的对齐方式:“RIGHT_LEFT”(默认)、“LEFT_RIGHT”、“LEFT_LEFT” 和 “RIGHT_RIGHT”。“RIGHT_LEFT” 将超对角线向右对齐(左侧填充行),将次对角线向左对齐(右侧填充行)。这是 LAPACK 使用的打包格式。cuSPARSE 使用 “LEFT_RIGHT”,这是相反的对齐方式。

返回:

result – 具有给定对角线值的新 tensor。

返回类型:

tvm.te.Tensor

示例

data = [[[7, 7, 7, 7],
         [7, 7, 7, 7],
         [7, 7, 7, 7]],
        [[7, 7, 7, 7],
         [7, 7, 7, 7],
         [7, 7, 7, 7]]]

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

topi.matrix_set_diag(input, diagonal) =
    [[[1, 7, 7, 7],
      [7, 2, 7, 7],
      [7, 7, 3, 7]],
     [[4, 7, 7, 7],
      [7, 5, 7, 7],
      [7, 7, 6, 7]]]
tvm.topi.max(data, axis=None, keepdims=False)

给定轴或轴列表的数组元素的最大值

参数:
  • data (tvm.te.Tensor) – 输入 tvm 张量

  • axis (Noneinttuple 类型的 int) – 执行 max 操作的轴或轴。默认值 axis=None 将找到输入数组所有元素中的最大元素。如果 axis 为负数,则从最后一个轴计数到第一个轴。

  • keepdims (bool) – 如果设置为 True,则被缩减的轴将作为大小为 1 的维度保留在结果中。使用此选项,结果将与输入数组正确广播。

返回:

ret

返回类型:

tvm.te.Tensor

tvm.topi.maximum(lhs, rhs)

取两个张量的逐元素最大值,带自动广播

参数:
返回:

ret – 如果两个操作数都是 Expr,则返回 Expr。否则返回 Tensor。

返回类型:

tvm.te.Tensor 或 Expr

tvm.topi.meshgrid(a_tuple, indexing)

从坐标向量创建坐标矩阵。

参数:
  • a_tuple (tuple 类型的 tvm.te.Tensor) – 坐标向量或标量。

  • indexing (str) – 索引模式,可以是 “ij” 或 “xy”。

返回:

result – 每个轴的结果网格。

返回类型:

tuple 类型的 tvm.te.Tensor

tvm.topi.min(data, axis=None, keepdims=False)

给定轴或轴列表的数组元素的最小值

参数:
  • data (tvm.te.Tensor) – 输入 tvm 张量

  • axis (Noneinttuple 类型的 int) – 执行最小值操作的轴或轴。默认值 axis=None 将找到输入数组所有元素中的最小元素。如果 axis 为负数,则从最后一个轴计数到第一个轴。

  • keepdims (bool) – 如果设置为 True,则被缩减的轴将作为大小为 1 的维度保留在结果中。使用此选项,结果将与输入数组正确广播。

返回:

ret

返回类型:

tvm.te.Tensor

tvm.topi.minimum(lhs, rhs)

取两个张量的逐元素最大值,带自动广播

参数:
返回:

ret – 如果两个操作数都是 Expr,则返回 Expr。否则返回 Tensor。

返回类型:

tvm.te.Tensor 或 Expr

tvm.topi.mod(lhs, rhs)

自动广播取模

参数:
返回:

ret – 如果两个操作数都是 Expr,则返回 Expr。否则返回 Tensor。

返回类型:

tvm.te.Tensor 或 Expr

tvm.topi.multiply(lhs, rhs)

自动广播乘法

参数:
返回:

ret – 如果两个操作数都是 Expr,则返回 Expr。否则返回 Tensor。

返回类型:

tvm.te.Tensor 或 Expr

tvm.topi.ndarray_size(array, dtype='int32')

获取输入数组的元素数量

参数:
  • array (tvm.te.Tensor) – 源 tensor。

  • dtype (str, optional) – 目标数据类型。

返回:

result – 结果张量。

返回类型:

tvm.te.Tensor

tvm.topi.negative(x)

取输入 x 的负数。

参数:

x (tvm.te.Tensor) – 输入参数。

返回:

y – 结果。

返回类型:

tvm.te.Tensor

tvm.topi.not_equal(lhs, rhs)

使用自动广播计算 (lhs!=rhs)

参数:
返回:

ret – 如果两个操作数都是 Expr,则返回 Expr。否则返回 Tensor。

返回类型:

tvm.te.Tensor 或 Expr

tvm.topi.one_hot(indices, on_value, off_value, depth, axis, dtype)

返回一个独热 tensor,其中由 indices 表示的位置取值 on_value,其他位置取值 off_value。最终维度是 <indices 外层维度> x depth x <indices 内层维度>。

参数:
  • indices (tvm.te.Tensor) – 要设置为 on_value 的位置。

  • on_value (tvm.te.Tensor) – 在 indices 处填充的值。

  • off_value (tvm.te.Tensor) – 在 indices 之外的所有其他位置填充的值。

  • depth (int) – 独热维度的深度。

  • axis (int) – 要填充的轴。

  • dtype (str) – 输出 tensor 的数据类型。

返回:

ret – 独热 tensor。

返回类型:

tvm.te.Tensor

示例

indices = [0, 1, 2]

topi.one_hot(indices, 3) =
    [[1, 0, 0],
     [0, 1, 0],
     [0, 0, 1]]
tvm.topi.power(lhs, rhs)

自动广播幂运算

参数:
返回:

ret – 如果两个操作数都是 Expr,则返回 Expr。否则返回 Tensor。

返回类型:

tvm.te.Tensor 或 Expr

tvm.topi.prod(data, axis=None, keepdims=False)

给定轴或轴列表的数组元素的乘积

参数:
  • data (tvm.te.Tensor) – 输入 tvm 张量

  • axis (Noneinttuple 类型的 int) – 执行 prod 操作的轴或轴。默认值 axis=None 将获得输入数组所有元素的乘积。如果 axis 为负数,则从最后一个轴计数到第一个轴。

  • keepdims (bool) – 如果设置为 True,则被缩减的轴将作为大小为 1 的维度保留在结果中。使用此选项,结果将与输入数组正确广播。

返回:

ret

返回类型:

tvm.te.Tensor

tvm.topi.reinterpret(x, dtype)

将输入重新解释为指定的数据类型。

参数:
返回:

y – 结果。

返回类型:

tvm.te.Tensor

tvm.topi.repeat(a, repeats, axis)

重复数组的元素。

参数:
  • a (tvm.te.Tensor) – 要重复的 tensor。

  • repeats (int,必需) – 每个元素的重复次数

  • axis (int,可选) – 要沿其重复值的轴

返回:

ret

返回类型:

tvm.te.Tensor

tvm.topi.reshape(a, newshape)

重塑数组

参数:
  • a (tvm.te.Tensor) – 要改变形状的 tensor

  • newshape (tuple 类型的 ints) – 新形状

返回:

ret

返回类型:

tvm.te.Tensor

tvm.topi.reverse_sequence(a, seq_lengths, seq_axis=1, batch_axis=0)

反转可变长度切片的 tensor。首先沿批次轴对输入进行切片,然后沿序列轴反转元素。

参数:
  • a (tvm.te.Tensor) – 要反转的 tensor。

  • seq_lengths (tvm.te.Tensor) – 长度为 a.dims[batch_axis] 的 1D Tensor。必须是以下类型之一:int32、int64。如果 seq_lengths[i] > a.dims[seq_axis],则四舍五入为 a.dims[seq_axis]。如果 seq_lengths[i] < 1,则四舍五入为 1。

  • seq_axis (int,可选) – 元素将沿其反转的轴。默认为 1。

  • batch_axis (int,可选) – tensor 将沿其切片的轴。默认为 0。

返回:

ret – 与输入形状和类型相同的计算结果。

返回类型:

tvm.te.Tensor

tvm.topi.right_shift(lhs, rhs)

自动广播右移

参数:
返回:

ret – 如果两个操作数都是 Expr,则返回 Expr。否则返回 Tensor。

返回类型:

tvm.te.Tensor 或 Expr

tvm.topi.round(x)

将 x 的元素四舍五入到最接近的整数。

参数:

x (tvm.te.Tensor) – 输入参数。

返回:

y – 结果。

返回类型:

tvm.te.Tensor

tvm.topi.rsqrt(x)

取输入 x 的平方根倒数。

参数:

x (tvm.te.Tensor) – 输入参数。

返回:

y – 结果。

返回类型:

tvm.te.Tensor

tvm.topi.scanop(data: Tensor, binop: Callable[[tvm.Expr, tvm.Expr], tvm.Expr], identity_value: tvm.Expr, op_name: str, axis: int | None = None, dtype: str | None = None, exclusive: bool | None = None) Tensor

累积二元运算符(扫描),其轴行为类似于 np.cumsum 和 np.cumprod。

有关使用示例,请参见 cumprod 和 cumsum。

例如,如果 * 是您的二元运算符,并且输入 tensor 是 [1, 2, 3, 4],则输出可能是 [1, 1 * 2, 1 * 2 * 3, 1 * 2 * 3 * 4]

参数:
  • data (tvm.te.Tensor) – 运算符的输入数据。

  • binop (Callable (tvm.Expr, tvm.Expr) -> tvm.Expr) – 应该是关联的和可交换的二元运算符。例如,如果 * 是您的运算符,则 a * (b * c) = (a * b) * c 且 a * b = b * a

  • identity_value (tvm.Expr) – 二元运算的值,提供恒等属性。例如,如果 * 是您的运算符,而 i 是 identity_value,则对于运算域中的所有 a,a * i = a。

  • axis (int,可选) – 计算操作所沿的轴。默认值 (None) 是在展平的数组上计算累积操作。

  • dtype (string,可选) – 返回数组和在其中计算元素的累加器类型。如果未指定 dtype,则默认为 data 的 dtype。

  • exclusive (bool,可选) – 如果为 True,将返回排除性的累积操作,其中不包括第一个元素。换句话说,如果为 True,则第 j 个输出元素将是前 (j-1) 个元素的累积操作。否则,它将是前 j 个元素的累积操作。零元素的累积操作假定为 identity_value。

返回:

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

返回类型:

tvm.te.Tensor

tvm.topi.scatter_elements(data, indices, updates, axis=0, reduction='update')

将更新中的元素分散到复制数据的相应索引中。

Data、indices、updates 和输出具有相同的形状。如果 reduction == “update”,则 Indices 不能有重复项(如果 idx1 != idx2,则 indices[idx1] != indices[idx2])。

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 由 reduction 确定。支持五种类型的函数:“update”、“add”、“mul”、“min” 和 “max”(见下文)

参数:
  • data (tvm.te.Tensor) – 源数组。

  • indices (tvm.te.Tensor) – 要提取的值的索引。

  • updates (tvm.te.Tensor) – 要在 Indices 处应用的更新

  • axis (optional, int) – 要在其上进行散布的轴。默认为零。

  • reduction (optional, string) – 算法的更新模式,可以是 “update”、“add”、“mul”、“min” 或 “max”。如果为 update,则更新值将替换输入数据。如果为 add,则更新值将添加到输入数据。如果为 mul,则输入数据将乘以更新值。如果为 mean,则输入数据将是更新值和输入数据之间的平均值。如果为 min,则在更新值和输入数据之间选择最小值。如果为 max,则在更新值和输入数据之间选择最大值。默认为 “update”

返回:

ret

返回类型:

tvm.te.Tensor

tvm.topi.scatter_nd(data, indices, updates, mode)

从 n 维数组中分散元素。

给定形状为 (Y_0, …, Y_{K-1}, X_M, …, X_{N-1}) 的 updates,形状为 (M, Y_0, …, Y_{K-1}) 的 indices,以及从形状为 (X_0, X_1, …, X_{N-1}) 的 data 复制的输出,scatter_nd 计算

output[indices[0, y_0, ..., y_{K-1}],
       ...,
       indices[M-1, y_0, ..., y_{K-1}],
       x_M,
       ...,
       x_{N-1}
      ] = f(output[...], updates[y_0, ..., y_{K-1}, x_M, ..., x_{N-1}])

其中更新函数 f 由 mode 确定。

参数:
  • data (tvm.te.Tensor) – 源数组。

  • indices (tvm.te.Tensor) – 要提取的值的索引。

  • updates (tvm.te.Tensor) – 要在 Indices 处应用的更新

  • mode (string) – 算法的更新模式,可以是 “update” 或 “add”。如果为 update,则更新值将替换输入数据。如果为 add,则更新值将添加到输入数据。

返回:

ret

返回类型:

tvm.te.Tensor

tvm.topi.searchsorted(sorted_sequence, values, right=False, out_dtype='int64')
查找应插入元素以保持顺序的索引。

如果 sorted_sequence 是 N 维的,则在 sorted_sequence 的相应维度中搜索 values 的最内层维度。

参数:
  • sorted_sequence (te.Tensor) – N 维或 1 维 Tensor,在最内层维度上包含单调递增的序列。

  • values (te.Tensor) – 包含搜索值的 N 维 Tensor。当 sorted_sequence 为 1 维时,values 的形状可以是任意的。否则,sorted_sequencevalues 的秩必须相同,并且外层 N-1 轴必须具有相同的大小。

  • right (bool,可选) – 控制当值完全落在已排序值之一上时返回哪个索引。如果为 False,则给出找到的第一个合适位置的索引。如果为 true,则返回最后一个此类索引。如果没有合适的索引,则返回 0 或 N(其中 N 是最内层维度的大小)。

  • dtype (string,可选) – 输出索引的数据类型。

返回:

indices – 形状与 values 相同的 Tensor,表示 values 的元素如果插入到 sorted_sequence 中的索引。

返回类型:

te.Tensor

tvm.topi.sequence_mask(data, valid_length, mask_value=0, axis=0)

将序列预期长度之外的所有元素设置为常量值。

此函数采用 n 维输入数组,形式为 [MAX_LENGTH, batch_size, …] 或 [batch_size, MAX_LENGTH, …],并返回形状相同的数组。

axis 表示长度维度的轴,只能为 0 或 1。如果 axis 为 0,则数据必须具有形状 [MAX_LENGTH, batch_size, …]。否则 (axis=1),数据必须具有形状 [batch_size, MAX_LENGTH, …]。

valid_length 给出了每个序列的长度。valid_length 应该是一个 1D int 数组,包含正整数,维度为 [batch_size,]。

参数:
  • data (tvm.te.Tensor) – N 维,形状为 [MAX_LENGTH, batch_size, …] 或 [batch_size, MAX_LENGTH, …],具体取决于 axis 的值。

  • valid_length (tvm.te.Tensor) – 1 维,形状为 [batch_size,]

  • mask_value (float,可选) – 掩码值,默认为 0

  • axis (int,可选) – 长度维度的轴,必须为 0 或 1,默认为 0

返回:

output – N 维,形状为 [MAX_LENGTH, batch_size, …] 或 [batch_size, MAX_LENGTH, …],具体取决于 axis 的值。

返回类型:

tvm.te.Tensor

tvm.topi.shape(array, dtype='int32')

获取输入数组的形状

参数:
  • array (tvm.te.Tensor) – 源 tensor。

  • dtype (str, optional) – 目标数据类型。

返回:

result – 结果张量。

返回类型:

tvm.te.Tensor

tvm.topi.sigmoid(x)

取输入 x 的 sigmoid tanh。

参数:

x (tvm.te.Tensor) – 输入参数。

返回:

y – 结果。

返回类型:

tvm.te.Tensor

tvm.topi.sign(x)

根据 x 的符号返回 -1、0、1。

参数:

x (tvm.te.Tensor) – 输入参数。

返回:

y – 结果。

返回类型:

tvm.te.Tensor

tvm.topi.sin(x)

取输入 x 的正弦。

参数:

x (tvm.te.Tensor) – 输入参数。

返回:

y – 结果。

返回类型:

tvm.te.Tensor

tvm.topi.sinh(x)

取输入 x 的双曲正弦。

参数:

x (tvm.te.Tensor) – 输入参数。

返回:

y – 结果。

返回类型:

tvm.te.Tensor

tvm.topi.sliding_window(data, axis, window_shape, strides)

在数据张量上滑动窗口。

参数:
  • data (tvm.te.Tensor) – 运算符的输入数据。

  • axis (int) – 窗口开始在其上滑动的轴。窗口将在此轴和所有后续轴上滑动。轴值确定窗口形状(以及步幅数):窗口形状和步幅的长度都必须为 data.ndim-axis

  • window_shape (List[int]) – 要在输入上形成的窗口形状。窗口形状的长度必须为 data.ndim-axis

  • strides (List[int]) – 如何沿每个维度步进窗口。步幅的长度必须为 data.ndim-axis

返回:

result – 结果张量。

返回类型:

tvm.te.Tensor

tvm.topi.sort(data, axis=-1, is_ascend=1)

沿给定轴执行排序,并返回排序顺序的数组。

参数:
  • data (tvm.te.Tensor) – 输入张量。

  • axis (int, optional) – 对输入张量进行排序的轴。默认情况下使用展平的数组。

  • is_ascend (boolean, optional) – 是否按升序或降序排序。

  • dtype (string, optional) – 输出索引的数据类型。

返回:

out – 排序后的索引张量。

返回类型:

tvm.te.Tensor

tvm.topi.sparse_reshape(sparse_indices, prev_shape, new_shape, new_sparse_indices_shape, new_shape_shape)

重塑稀疏张量

参数:
  • sparse_indices (te.Expr) – 一个 2-D tensor[N, n_dim],包含稀疏值的位置,其中 N 是稀疏值的数量,n_dim 是 dense_shape 的维度数

  • prev_shape (te.Expr) – 一个 1-D tensor,包含稠密 tensor 的先前形状

  • new_shape (te.Expr) – 一个 1-D tensor,包含稠密 tensor 的新形状

返回:

result – 输出 tensor。

返回类型:

te.Expr

示例

sparse_indices = [[0, 0, 0],
                    [0, 0, 1],
                    [0, 1, 0],
                    [1, 0, 0],
                    [1, 2, 3]]
prev_shape = [2, 3, 4]
new_shape = [9, -1]
new_sparse_indices, new_shape = topi.sparse_reshape(
    sparse_indices, prev_shape, new_shape)
new_sparse_indices = [[0, 0],
                      [0, 1],
                      [1, 2],
                      [4, 2],
                      [8, 1]]
new_shape = [9, 4]
tvm.topi.sparse_to_dense(sparse_indices, output_shape, sparse_values, default_value=0)

将稀疏表示形式转换为稠密张量。

示例:: - sparse_to_dense([[0, 0], [1, 1]], [2, 2], [3, 3], 0) = [[3, 0], [0, 3]]

参数:
  • sparse_indices (tvm.te.Tensor) – 一个 0-D、1-D 或 2-D tensor,包含稀疏值的位置。

  • output_shape (整数列表) – 稠密输出 tensor 的形状。

  • sparse_values (tvm.te.Tensor) – 一个 0-D 或 1-D tensor,包含稀疏索引的稀疏值。

  • default_value (tvm.te.Tensor) – 一个 0-D tensor,包含剩余位置的默认值。默认为 0。

返回:

result – 形状为 output_shape 的稠密 tensor。与 sparse_values 具有相同的类型。

返回类型:

tvm.te.Tensor

tvm.topi.split(ary, indices_or_sections, axis=0)

将数组拆分为多个子数组。

参数:
返回:

ret

返回类型:

tuple 类型的 tvm.te.Tensor

tvm.topi.sqrt(x)

取输入 x 的平方根。

参数:

x (tvm.te.Tensor) – 输入参数。

返回:

y – 结果。

返回类型:

tvm.te.Tensor

tvm.topi.squeeze(a, axis=None)

从数组的形状中删除单维度条目。

参数:
  • a (tvm.te.Tensor)

  • axis (Noneinttuple 类型的 ints,可选) – 选择形状中单维度条目的子集。如果选择的轴的形状条目大于 1,则会引发错误。

返回:

挤压后的

返回类型:

tvm.te.Tensor

tvm.topi.stack(a, axis)

多次重复整个数组。

参数:
  • a (tvm.te.Tensor) – 要堆叠的 tensor。

  • axis (int,可选) – 输入数组沿其堆叠的结果数组中的轴。

返回:

ret

返回类型:

tvm.te.Tensor

tvm.topi.stft(data, n_fft, hop_length, win_length, window, normalized, onesided, output_shape)

STFT 计算输入信号短时重叠窗口的傅里叶变换。这给出了信号的频率分量随时间的变化。 :param data: 1-D 张量或 2-D 批次张量。 :type data: te.Tensor :param n_fft: 傅里叶变换的大小 :type n_fft: int :param hop_length: 相邻滑动窗口帧之间的距离 :type hop_length: int :param win_length: 窗口帧和 STFT 滤波器的大小 :type win_length: int :param window: 1-D 张量窗口帧 :type window: te.Tensor :param normalized: 是否返回归一化的 STFT 结果 :type normalized: bool :param onesided: 是否返回单边结果或填充共轭对称 :type onesided: bool

返回:

output – 包含 STFT 结果的张量

返回类型:

te.Tensor

示例

data = [1, 2, 3, 4, 5, 6]
window = [4, 3, 2]
[n_fft, hop_length, win_length, normalized, onesided] = [3, 3, 3, False, True]
topi.stft(data, n_fft, hop_length, win_length, window, normalized, onesided)
-> [[[15.0000,  0.0000], [34.0000,  0.0000]], [[ 4.5000,  0.8660], [ 1.0000, -1.7321]]]
tvm.topi.strided_set(a, v, begin, end, strides=None)

设置数组切片。

参数:
返回:

ret

返回类型:

tvm.te.Tensor

tvm.topi.strided_slice(a, begin, end, strides=None, axes=None, slice_mode='end', assume_inbound=True)

数组切片。

参数:
  • a (tvm.te.Tensor) – 要切片的张量。

  • begin (list of int) – 切片开始的索引。

  • end (list of int) – 指示切片结束的索引。

  • strides (list of int, optional) – 指定步长值,可以是负数,在这种情况下,输入张量将在该特定轴上反转。

  • axes (list of int, optional) – 应用切片的轴。当指定时,begin、end、strides 和 axes 需要是相同长度的整数列表。

  • slice_mode (str, optional) – 切片模式 [end, size]。end - 切片的结束索引 [默认]。size - 输入步长将被忽略,此模式下的输入 end 指示从 begin 指定的位置开始的切片的大小。如果 end[i] 为 -1,则该维度中的所有剩余元素都包含在切片中。

  • assume_inbound (bool, optional) – 一个标志,指示是否假定所有索引都在界内

返回:

ret

返回类型:

tvm.te.Tensor

tvm.topi.subtract(lhs, rhs)

自动广播减法

参数:
返回:

ret – 如果两个操作数都是 Expr,则返回 Expr。否则返回 Tensor。

返回类型:

tvm.te.Tensor 或 Expr

tvm.topi.sum(data, axis=None, keepdims=False)

给定轴或轴列表的数组元素的总和

参数:
  • data (tvm.te.Tensor) – 输入 tvm 张量

  • axis (None or int or tuple of int) – 执行求和运算的轴或轴。默认值 axis=None 将对输入数组的所有元素求和。如果 axis 为负数,则从最后一个轴计数到第一个轴。

  • keepdims (bool) – 如果设置为 True,则被缩减的轴将作为大小为 1 的维度保留在结果中。使用此选项,结果将与输入数组正确广播。

返回:

ret

返回类型:

tvm.te.Tensor

tvm.topi.take(a, indices, axis=None, batch_dims=0, mode='clip')

沿轴从数组中获取元素。

参数:
  • a (tvm.te.Tensor) – 源数组。

  • indices (tvm.te.Tensor) – 要提取的值的索引。

  • axis (int, optional) – 选择值的轴。默认情况下,使用展平的输入数组。

  • batch_dims (int) – 批次维度的数量。默认为 0。

  • mode (str, optional) – 指定越界索引的行为。clip - 裁剪到范围(默认) wrap - 环绕索引 fast - 无裁剪或环绕(用户必须确保索引在界内)

返回:

ret

返回类型:

tvm.te.Tensor

tvm.topi.tan(x)

取输入 x 的正切。

参数:

x (tvm.te.Tensor) – 输入参数。

返回:

y – 结果。

返回类型:

tvm.te.Tensor

tvm.topi.tanh(x)

取输入 x 的双曲正切。

参数:

x (tvm.te.Tensor) – 输入参数。

返回:

y – 结果。

返回类型:

tvm.te.Tensor

tvm.topi.tensordot(a, b, axes)

矩阵乘法到张量的推广。

参数:
  • a (张量 A)

  • b (张量 B)

  • axes (要缩减的维度数量)

返回类型:

计算结果的张量

tvm.topi.tile(a, reps)

多次重复整个数组。

参数:
  • a (tvm.te.Tensor) – 要平铺的张量。

  • reps (tuple of ints, required) – 重复张量的次数

返回:

ret

返回类型:

tvm.te.Tensor

tvm.topi.topk(data, k=1, axis=-1, ret_type='both', is_ascend=False, dtype='int64')

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

参数:
  • data (tvm.te.Tensor) – 输入张量。

  • k (int or tvm.te.Tensor, optional) – 要选择的顶部元素的数量。如果 k < 1,则返回所有元素。

  • axis (int, optional) – 对输入张量进行排序的轴。

  • ret_type (str, optional) – 返回类型 [both, values, indices]。 “both”:返回顶部 k 个数据和索引。“values”:仅返回顶部 k 个数据。“indices”:仅返回顶部 k 个索引。

  • is_ascend (boolean, optional) – 是否按升序或降序排序。

  • dtype (string, optional) – 索引输出的数据类型。

返回:

out – 计算结果。

返回类型:

tvm.te.Tensor 或 List[tvm.te.Tensor]

tvm.topi.transpose(a, axes=None)

置换数组的维度。

参数:
  • a (tvm.te.Tensor) – 要扩展的张量。

  • axes (tuple of ints, optional) – 默认情况下,反转维度。

返回:

ret

返回类型:

tvm.te.Tensor

tvm.topi.trilu(data, k, upper)

给定 2-D 矩阵或批量的 2-D 矩阵,返回张量的上或下三角部分。

参数:
  • data (tvm.te.Tensor) – 将应用 trilu 的张量。必须是 2D 矩阵或 2D 矩阵批次的张量。

  • k (tvm.te.Tensor) – 要排除或包括在主对角线上方或下方的对角线数量。

  • upper (bool) – 如果为 True,则仅保留输入的上三角值;如果为 False,则保留下三角值。

返回:

ret – 对角线设置为零的新张量。

返回类型:

tvm.te.Tensor

示例

x = [[0, 1, 2],
     [3, 4, 5],
     [6, 7, 8]]

topi.trilu(x, True, 0) =
    [[0, 1, 2],
     [0, 4, 5],
     [0, 0, 8]]
tvm.topi.trunc(x)

逐元素取输入 x 的截断值。

参数:

x (tvm.te.Tensor) – 输入参数。

返回:

y – 结果。

返回类型:

tvm.te.Tensor

tvm.topi.unravel_index(indices, shape)

将平面索引或平面索引数组转换为坐标数组元组。

Example:: - unravel_index([22, 41, 37], [7, 6]) = [[3, 6, 6], [4, 5, 1]]

参数:
返回:

result – 坐标数组的元组。

返回类型:

tvm.te.Tensor

tvm.topi.where(condition, x, y)

根据条件从 x 或 y 中获取元素。

参数:
返回:

result – 根据条件从 x 或 y 中选择的张量。

返回类型:

tvm.te.Tensor

tvm.topi.within_index(b, e, s, i)

返回一个布尔值,指示 i 是否在给定索引内。

参数:
  • b (Expr) – 索引的起始位置

  • e (Expr) – 索引的结束位置

  • s (Expr) – 索引的步幅

  • i (Expr) – 数组位置

返回:

selected – 布尔表达式,如果数组位置将被索引选中则为 True,否则为 False

返回类型:

Expr

exception tvm.topi.InvalidShapeError

topi 函数的无效形状。例如,为非 3x3 内核调用 winograd 模板)

tvm.topi.nn

神经网络运算符

Workload(in_dtype, out_dtype, height, width, ...)

函数

adaptive_pool(data, output_size, pool_type)

对数据的高度和宽度维度执行池化。

adaptive_pool1d(data, output_size, pool_type)

对三维数据执行池化。

adaptive_pool3d(data, output_size, pool_type)

对三维数据执行池化。

add(lhs, rhs)

自动广播加法

batch_matmul(tensor_a, tensor_b[, oshape, ...])

计算 tensor_atensor_b 的批次矩阵乘法。

batch_norm(data, gamma, beta, moving_mean, ...)

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

batch_to_space_nd(data, block_shape, ...)

对数据执行空间到批次转换

binarize_pack(data[, axis, name])

沿特定轴进行二值化和位打包。

binary_dense(data, weight)

使用异或和位计数的二进制矩阵乘法。

bitpack(data, bits, pack_axis, bit_axis, ...)

将数据打包成位串行计算所需的格式

bitserial_conv2d_nchw(data, kernel, stride, ...)

位串行 Conv2D 运算符。

bitserial_conv2d_nhwc(data, kernel, stride, ...)

位串行 Conv2D 运算符。

bitserial_dense(data, weight, data_bits, ...)

topi 中位串行密集连接层的默认实现。

concatenate(a_tuple[, axis])

沿现有轴连接数组序列。

conv(inp, filt, stride, padding, dilation, ...)

NCHW 或 NHWC 布局中的卷积运算符。

conv1d(data, kernel[, strides, padding, ...])

1D 卷积前向运算符。

conv1d_ncw(data, kernel[, strides, padding, ...])

NCW 布局中的 1D 卷积。

conv1d_nwc(data, kernel[, strides, padding, ...])

NWC 布局中的 1D 卷积。

conv1d_transpose_ncw(data, kernel, stride, ...)

转置 1D 卷积 ncw 前向运算符。

conv2d(input, filter, strides, padding, dilation)

Conv2D 运算符。

conv2d_NCHWc(data, kernel, stride, padding, ...)

用于 nChw[x]c 布局的 Conv2D 运算符。

conv2d_NCHWc_int8(data, kernel, stride, ...)

用于 nChw[x]c 布局的 Conv2D 运算符。

conv2d_hwcn(Input, Filter, stride, padding, ...)

HWCN 布局中的卷积运算符。

conv2d_nchw(Input, Filter, stride, padding, ...)

NCHW 布局中的卷积运算符。

conv2d_nhwc(Input, Filter, stride, padding, ...)

NHWC 布局中的卷积运算符。

conv2d_transpose_nchw(Input, Filter, ...)

转置 2D 卷积 nchw 前向运算符。

conv2d_transpose_nchw_preprocess(data, ...)

预处理数据和内核,使 conv2d_transpose 的计算模式与 conv2d 相同

conv2d_winograd_nchw(data, weight, strides, ...)

NCHW 布局中的 Conv2D Winograd。

conv2d_winograd_nchw_without_weight_transform(...)

NCHW 布局中没有布局转换的 Conv2D Winograd。

conv2d_winograd_nhwc(data, weight, strides, ...)

NHWC 布局中的 Conv2D Winograd。

conv2d_winograd_nhwc_without_weight_transform(...)

NHWC 布局中没有布局转换的 Conv2D Winograd。

conv2d_winograd_weight_transform(kernel, ...)

winograd 的权重变换

conv3d_ncdhw(Input, Filter, stride, padding, ...)

NCDHW 布局中的 Conv3D 运算符。

conv3d_ndhwc(Input, Filter, stride, padding, ...)

NDHWC 布局中的卷积运算符。

conv3d_transpose_ncdhw(Input, Filter, ...)

转置 3D 卷积 ncdhw 前向运算符。

conv3d_transpose_ncdhw_preprocess(data, ...)

预处理数据和内核,使 conv3d_transpose 的计算模式与 conv3d 相同

conv3d_winograd_weight_transform(kernel, ...)

3D winograd 的权重变换

correlation_nchw(data1, data2, kernel_size, ...)

NCHW 布局中的相关运算符。

declaration_conv2d_transpose_impl(data, ...)

conv2d 转置的实现

declaration_conv3d_transpose_impl(data, ...)

conv3d 转置的实现

deformable_conv2d_nchw(data, offset, kernel, ...)

NCHW 布局中的可变形 conv2D 运算符。

deformable_conv2d_nhwc(data, offset, kernel, ...)

NHWC 布局中的可变形 conv2D 运算符。

dense(data, weight[, bias, out_dtype, ...])

topi 中密集连接层的默认实现。

dense_pack(data, weight[, bias, out_dtype])

topi 中 dense_pack 的默认实现。

depth_to_space(data, block_size[, layout, mode])

对数据执行深度到空间转换

depthwise_conv2d_NCHWc(Input, Filter, ...[, ...])

深度可分离卷积 NCHW[x]c 前向运算符。

depthwise_conv2d_backward_input_nhwc(Filter, ...)

深度可分离卷积 nhwc 反向传播输入运算符。

depthwise_conv2d_backward_weight_nhwc(Input, ...)

深度可分离卷积 nhwc 反向传播权重运算符。

depthwise_conv2d_nchw(Input, Filter, stride, ...)

深度可分离卷积 nchw 前向运算符。

depthwise_conv2d_nhwc(Input, Filter, stride, ...)

深度可分离卷积 nhwc 前向运算符。

dilate(data, strides[, dilation_value, name])

使用给定的膨胀值(默认为 0)膨胀数据。

equal_const_int(expr, value)

返回 expr 是否等于 value。

fast_softmax(x[, axis])

对数据执行 softmax 激活。

fifo_buffer(data, buffer, axis)

FIFO 缓冲区,用于在具有滑动窗口输入的 CNN 中实现计算重用

flatten(data)

通过折叠更高维度,将输入数组展平为 2-D 数组。

get_const_int(expr)

验证 expr 是否为整数并获取常量值。

get_const_tuple(in_tuple)

验证输入元组是否为 IntImm 或 Var,返回 int 或 Var 元组。

get_pad_tuple(padding, kernel)

获取填充选项的通用代码

get_pad_tuple1d(padding, kernel)

获取填充选项的通用代码

get_pad_tuple3d(padding, kernel)

获取填充选项的通用代码

get_pad_tuple_generic(padding, kernel)

获取填充选项的通用代码

global_pool(data, pool_type[, layout])

对数据的高度和宽度维度执行全局池化。

group_conv1d_ncw(data, kernel[, strides, ...])

NCW 布局的 1D 卷积前向运算符。

group_conv1d_nwc(data, kernel[, strides, ...])

NWC 布局的 1D 卷积前向运算符。

group_conv1d_transpose_ncw(data, kernel, ...)

转置 1D 分组卷积 ncw 前向运算符。

group_conv2d_nchw(Input, Filter, stride, ...)

NCHW 布局中的分组卷积运算符。

group_conv2d_nhwc(Input, Filter, stride, ...)

NHWC 布局中的分组卷积运算符。

group_conv2d_transpose_nchw(data, kernel, ...)

NCHW 布局中的分组卷积运算符。

group_conv3d_transpose_ncdhw(data, kernel, ...)

转置分组 3D 卷积 ncdhw 前向运算符。

leaky_relu(x, alpha)

对输入 x 执行 leaky relu。

log_softmax(x[, axis])

对数据执行 log softmax 激活

lrn(data, size[, axis, alpha, beta, bias])

对输入数据执行跨通道局部响应归一化。

lstm(Xs, Wi, Wh[, Bi, Bh, h_init, c_init, ...])

使用 TE scan 实现的通用 LSTM。

matmul(tensor_a, tensor_b[, bias, ...])

topi 中 matmul 的默认实现。

mirror_pad(data, pad_before[, pad_after, ...])

使用对称或反射镜像填充输入。

namedtuple(typename, field_names, *[, ...])

返回具有命名字段的元组的新子类。

nll_loss(predictions, targets, weights, ...)

输入数据的负对数似然损失。

pad(data, pad_before[, pad_after, ...])

用零填充输入。

pool1d(data, kernel, stride, dilation, ...)

对数据的宽度维度执行池化。

pool2d(data, kernel, stride, dilation, ...)

对数据的高度和宽度维度执行池化。

pool3d(data, kernel, stride, dilation, ...)

对数据的深度、高度和宽度维度执行池化操作。

pool_grad(grads, data, kernel, stride, ...)

数据的高度和宽度维度上的池化梯度。

prelu(x, slope[, axis])

PReLU。

reduce(function, sequence[, initial])

从左到右,将一个接受两个参数的函数累积地应用于序列的项目,从而将序列减少为单个值。

relu(x)

对输入 x 执行 ReLU 运算。

scale_shift_nchw(Input, Scale, Shift)

推理中的批归一化运算符。

scale_shift_nchwc(Input, Scale, Shift)

推理中的批归一化运算符。

scale_shift_nhwc(Input, Scale, Shift)

推理中的批归一化运算符。

simplify(expr)

如果表达式是 Expr,则简化表达式;如果表达式是 int,则直接返回。

simulated_dequantize(data, in_dtype[, ...])

模拟 QNN 反量化运算符,模仿 QNN 输出,但不更改数据类型。

simulated_quantize(data, out_dtype[, ...])

模拟 QNN 量化运算符,模仿 QNN 输出,但不更改数据类型。

softmax(x[, axis])

对数据执行 softmax 激活。

softmax_common(x, axis, use_fast_exp)

softmax 和 fast_softmax 的公共部分

space_to_batch_nd(data, block_shape, ...[, ...])

对数据执行批次到空间转换

space_to_depth(data, block_size[, layout])

对数据执行空间到深度转换

strided_slice(a, begin, end[, strides, ...])

数组切片。

unpack_NCHWc_to_nchw(packed_out, out_dtype)

将 conv2d_NCHWc 输出从 NCHWc 布局解包到 NCHW

upsampling(data, scale_h, scale_w[, layout, ...])

对数据执行上采样。

upsampling3d(data, scale_d, scale_h, scale_w)

对数据执行上采样。

winograd_transform_matrices(tile_size, ...)

计算 tile_size 的 A、B 和 G 变换矩阵,结果为 tvm.Expr

instance_norm(data, gamma, beta, axis[, epsilon])

实例归一化运算符。

layer_norm(data, gamma, beta, axis[, epsilon])

层归一化运算符。

group_norm(data, gamma, beta, num_groups, ...)

组归一化运算符。

rms_norm(data, weight, axis[, epsilon])

均方根归一化运算符。

class tvm.topi.nn.Workload(in_dtype, out_dtype, height, width, in_filter, out_filter, kernel_h, kernel_w, padt, padl, padb, padr, dilation_h, dilation_w, stride_h, stride_w)

属性

dilation_h

字段编号 12 的别名

dilation_w

字段编号 13 的别名

height

字段编号 2 的别名

in_dtype

字段编号 0 的别名

in_filter

字段编号 4 的别名

kernel_h

字段编号 6 的别名

kernel_w

字段编号 7 的别名

out_dtype

字段编号 1 的别名

out_filter

字段编号 5 的别名

padb

字段编号 10 的别名

padl

字段编号 9 的别名

padr

字段编号 11 的别名

padt

字段编号 8 的别名

stride_h

字段编号 14 的别名

stride_w

字段编号 15 的别名

width

字段编号 3 的别名

dilation_h

字段编号 12 的别名

dilation_w

字段编号 13 的别名

height

字段编号 2 的别名

in_dtype

字段编号 0 的别名

in_filter

字段编号 4 的别名

kernel_h

字段编号 6 的别名

kernel_w

字段编号 7 的别名

out_dtype

字段编号 1 的别名

out_filter

字段编号 5 的别名

padb

字段编号 10 的别名

padl

字段编号 9 的别名

padr

字段编号 11 的别名

padt

字段编号 8 的别名

stride_h

字段编号 14 的别名

stride_w

字段编号 15 的别名

width

字段编号 3 的别名

tvm.topi.nn.adaptive_pool(data, output_size, pool_type, layout='NCHW')
对数据的高度和宽度维度执行池化。

池化核和步幅大小是根据所需的输出大小自动选择的。它根据布局字符串决定高度和宽度维度,其中 ‘W’ 和 ‘H’ 分别表示宽度和高度。宽度和高度维度不能拆分。例如,NCHW、NCHW16c 等对于池化是有效的,而 NCHW16w、NCHW16h 则无效。有关布局字符串约定的更多信息,请参阅参数 layout

参数:
  • data (tvm.te.Tensor) – n 维,形状为布局

  • output_size (tuple of int) – 输出高度和宽度。

  • pool_type (str) – 池化类型,‘max’ 或 ‘avg’

  • layout (string) – 输入数据的布局。布局应由大写字母、小写字母和数字组成,其中大写字母表示维度,相应的小写字母和因子大小表示拆分维度。例如,NCHW16c 可以描述一个 5 维张量,形状为 [batch_size, channel, height, width, channel_block],其中 channel_block=16 是维度通道的拆分。

返回:

output – n 维,布局相同

返回类型:

tvm.te.Tensor

tvm.topi.nn.adaptive_pool1d(data, output_size, pool_type, layout='NCW')

对三维数据执行池化操作。有关详细信息,请参阅上面的二维版本。

tvm.topi.nn.adaptive_pool3d(data, output_size, pool_type, layout='NCDHW')

对三维数据执行池化操作。有关详细信息,请参阅上面的二维版本。

tvm.topi.nn.add(lhs, rhs)

自动广播加法

参数:
返回:

ret – 如果两个操作数都是 Expr,则返回 Expr。否则返回 Tensor。

返回类型:

tvm.te.Tensor 或 Expr

tvm.topi.nn.batch_matmul(tensor_a, tensor_b, oshape=None, out_dtype=None, transpose_a=False, transpose_b=True, auto_scheduler_rewritten_layout='', meta_schedule_original_shape=None)

计算 tensor_atensor_b 的批次矩阵乘法。

tensor_atensor_b 都可以转置。出于历史原因,默认情况下我们使用 NT 格式 (transpose_a=False, transpose_b=True)。

参数:
  • tensor_a (tvm.te.Tensor) – 3 维,形状为 [batch, M, K] 或 [batch, K, M]。

  • tensor_b (tvm.te.Tensor) – 3 维,形状为 [batch, K, N] 或 [batch, N, K]。

  • oshape (List[Optional]) – 计算的显式预期输出形状。在具有动态输入形状的情况下可能很有用。

  • out_dtype (Optional[str]) – 指定混合精度批矩阵乘法的输出数据类型。

  • transpose_a (Optional[bool] = False) – 第一个张量是否为转置格式。

  • transpose_b (Optional[bool] = True) – 第二个张量是否为转置格式。

  • auto_scheduler_rewritten_layout (Optional[str] = "") – 自动调度器的布局重写传递后的布局。

  • meta_schedule_original_shape (Optional[List[PrimExpr]] = None) – 张量的原始形状

返回:

output – 3 维,形状为 [batch, M, N]

返回类型:

tvm.te.Tensor

tvm.topi.nn.batch_norm(data: Tensor, gamma: Tensor, beta: Tensor, moving_mean: Tensor, moving_var: Tensor, axis: int | None = None, epsilon: float | None = None, center: bool | None = None, scale: bool | None = None, training: bool | None = None, momentum: float | None = None) List[Tensor]

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

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

参数:
  • data (tvm.te.Tensor) – 要进行批归一化的输入。

  • gamma (tvm.te.Tensor) – 要应用于归一化张量的比例因子。

  • beta (tvm.te.Tensor) – 要应用于归一化张量的偏移量。

  • moving_mean (tvm.te.Tensor) – 输入的运行平均值。

  • moving_var (tvm.te.Tensor) – 输入的运行方差。

  • axis (int, optional, default=1) – 指定应沿哪个形状轴进行归一化。

  • epsilon (float, optional, default=1e-5) – 添加到方差中的小浮点数,以避免被零除。

  • center (bool, optional, default=True) – 如果为 True,则将 beta 偏移量添加到归一化张量;如果为 False,则忽略 beta。

  • scale (bool, optional, defualt=True) – 如果为 True,则按 gamma 缩放归一化张量;如果为 False,则忽略 gamma。

  • training (bool, optional, defualt=False) – 指示是否处于训练模式。如果为 True,则更新 moving_mean 和 moving_var。

  • momentum (float, optional, default=0.1) – 用于 moving_mean 和 moving_var 更新的值。

返回:

  • output (list of tvm.te.Tensor) – 归一化数据,形状与输入相同

  • moving_mean (tvm.te.Tensor) – 输入的运行平均值。

  • moving_var (tvm.te.Tensor) – 输入的运行方差。

tvm.topi.nn.batch_to_space_nd(data, block_shape, crop_begin_list, crop_end_list)

对数据执行空间到批次转换

参数:
  • data (tvm.te.Tensor) – N 维张量,形状为 [batch, spatial_shape, remaining_shapes],其中 spatial_shape 具有 M 个维度。

  • block_size (list of ints) – 大小为 [M] 的列表,其中 M 是空间维度数,指定每个空间维度的块大小。

  • crop_begin_list (list of ints) – 形状为 [M] 的列表,其中 M 是空间维度数,指定每个空间维度的起始裁剪大小。

  • crop_end_list (list of ints) – 形状为 [M] 的列表,其中 M 是空间维度数,指定每个空间维度的结束裁剪大小。

返回:

output

返回类型:

tvm.te.Tensor

tvm.topi.nn.binarize_pack(data, axis=None, name='PackedInput')

沿特定轴进行二值化和位打包。

参数:
  • data (tvm.te.Tensor) – n 维输入,可以是任何布局。

  • axis (None or int) – 执行二值化和位打包的轴,默认为最后一个轴。

  • name (str, optional) – 运算符生成的名字前缀。

返回:

output – n 维,与输入布局相同,dtype 为 uint32。

返回类型:

tvm.te.Tensor

tvm.topi.nn.binary_dense(data, weight)

使用异或和位计数的二进制矩阵乘法。

参数:
  • data (tvm.te.Tensor) – 2 维,形状为 [batch, in_dim],dtype 为 uint32。

  • weight (tvm.te.Tensor) – 2 维,形状为 [out_dim, in_dim],dtype 为 uint32。

返回:

output – 2 维,形状为 [batch, out_dim],dtype 为 float32。

返回类型:

tvm.te.Tensor

tvm.topi.nn.bitpack(data, bits, pack_axis, bit_axis, pack_type, name='QuantizeInput')

将数据打包成位串行计算所需的格式

参数:
  • pack_axis (int) – 要在数据中打包的轴的索引

  • bit_axis (int) – 在结果打包数据中放置位轴的轴的索引

tvm.topi.nn.bitserial_conv2d_nchw(data, kernel, stride, padding, activation_bits, weight_bits, pack_dtype='uint32', out_dtype='int16', unipolar=True)

位串行 Conv2D 运算符。

参数:
  • data (tvm.te.Tensor) – 4 维,形状为 [batch, in_channel, in_height, in_width]

  • kernel (tvm.te.Tensor) – 4 维,形状为 [num_filter, in_channel, filter_height, filter_width]

  • stride (int or a list/tuple of two ints) – 步幅大小,或 [stride_height, stride_width]

  • padding (int or a list/tuple of two or four ints) – 填充大小,[pad_height, pad_width], [pad_top, pad_left, pad_down, pad_right]

  • activation_bits (int) – 用于激活/输入元素的位数

  • weight_bits (int) – 用于权重元素的位数

  • out_dtype (str) – 卷积的返回类型

  • pack_dtype (str) – 位打包类型

  • unipolar (bool) – 如果二值化风格是单极性 1/0 格式,而不是双极性 -1/+1 格式

返回:

output – 4 维,形状为 [batch, out_channel, out_height, out_width]

返回类型:

tvm.te.Tensor

tvm.topi.nn.bitserial_conv2d_nhwc(data, kernel, stride, padding, activation_bits, weight_bits, pack_dtype='uint32', out_dtype='int16', unipolar=True)

位串行 Conv2D 运算符。

参数:
  • data (tvm.te.Tensor) – 4 维,形状为 [batch, in_height, in_width, in_channel]

  • kernel (tvm.te.Tensor) – 4 维,形状为 [filter_height, filter_width, in_channel, num_filter]

  • stride (int or a list/tuple of two ints) – 步幅大小,或 [stride_height, stride_width]

  • padding (int or a list/tuple of two or four ints) – 填充大小,[pad_height, pad_width], [pad_top, pad_left, pad_down, pad_right]

  • activation_bits (int) – 用于激活/输入元素的位数

  • weight_bits (int) – 用于权重元素的位数

  • out_dtype (str) – 卷积的返回类型

  • pack_dtype (str) – 位打包类型

  • unipolar (bool) – 如果二值化风格是单极性 1/0 格式,而不是双极性 -1/+1 格式

返回:

output – 4 维,形状为 [batch, out_height, out_width, out_channel]

返回类型:

tvm.te.Tensor

tvm.topi.nn.bitserial_dense(data, weight, data_bits, weight_bits, pack_dtype='uint32', out_dtype='int16', unipolar=True)

topi 中位串行密集连接层的默认实现。

参数:
  • data (tvm.te.Tensor) – 2 维,形状为 [batch, in_dim]

  • weight (tvm.te.Tensor) – 2 维,形状为 [out_dim, in_dim] 或 3 维,形状为 [out_dim, weight_bits, in_dim]

返回:

output – 2 维,形状为 [batch, out_dim]

返回类型:

tvm.te.Tensor

tvm.topi.nn.concatenate(a_tuple, axis=0)

沿现有轴连接数组序列。

参数:
  • a_tuple (tuple of tvm.te.Tensor) – 要连接的数组

  • axis (int, optional) – 数组将沿其连接的轴。默认为 0。

返回:

ret

返回类型:

tvm.te.Tensor

tvm.topi.nn.conv(inp: Tensor, filt: Tensor, stride: int | Sequence[int], padding: int | Sequence[int], dilation: int | Sequence[int], groups: int, data_layout: str, kernel_layout: str = '', out_dtype: str | None >= None, auto_scheduler_rewritten_layout: str | None >= None, meta_schedule_original_shape=None, auto_scheduler_should_rewrite_layout: bool >= False)

NCHW 或 NHWC 布局中的卷积运算符。

支持 1D、2D、3D 卷积以及分组卷积。

参数:
  • inp (tvm.te.Tensor) – N 维张量,形状为 [batch, in_channel, in_height, in_width, …],布局由 data_layout 指定

  • filt (tvm.te.Tensor) – N 维张量,形状为 [num_filter, in_channel // groups, filter_height, filter_width, …],布局由 kernel_layout 指定

  • stride (intdim 个整数的列表/元组) – (其中 NCHW 的 dim=2,NCH 的 dim=1,以此类推) 步幅大小,或 [stride_height, stride_width, …]

  • padding (intdim 或 2*dim 个整数的列表/元组) – (其中 NCHW 的 dim=2,NCH 的 dim=1,以此类推) 填充大小,或 dim 个整数的 [pad_height, pad_width, …],或 2*dim 个整数的 [pad_top, pad_left, pad_bottom, pad_right]

  • dilation (int两个整数的列表/元组) – 空洞大小,或 [dilation_height, dilation_width]

  • groups (int) – 分组数

  • data_layout (str) – 输入数据的布局。N 表示批次维度,C 表示通道,任何其他字符表示 HW(或 1D 的 H 和 3D 的 HWD)。

  • kernel_layout (Optional[str]) – 卷积核的布局。I 表示输入通道,O 表示输出通道,任何其他字符表示卷积核的 HW 维度(或 1D 的 H 和 3D 的 HWD)。如果 kernel_layout 为空,则使用 data_layout 推断默认的 kernel_layout。对于 NCHW 数据布局,默认的 kernel_layout 是 OIHW;对于 NHWC 数据布局,默认的 kernel_layout 是 HWIO。

  • out_dtype (str) – 元素在逐元素乘法和求和之前转换为此类型。

  • auto_scheduler_rewritten_layout (str) – 来自自动调度器的布局重写。

  • meta_schedule_original_shape (Optional[List[PrimExpr]]) – 输入张量的原始形状。

  • auto_scheduler_should_rewrite_layout (bool) – 是否允许自动调度器重写卷积核张量的布局。默认为 false。如果与分组卷积一起使用,可能会导致错误。

返回:

Output – N 维张量,形状为 [batch, out_channel, out_height, out_width, …],布局由 data_layout 指定

返回类型:

tvm.te.Tensor

tvm.topi.nn.conv1d(data, kernel, strides=1, padding='VALID', dilation=1, groups=1, data_layout='NCW', kernel_layout='', out_dtype=None)

1D 卷积前向运算符。

参数:
  • data (tvm.te.Tensor) – 3-D 输入,当 data_layout == ‘NCW’ 时形状为 [batch, in_channel, in_width],当 data_layout == ‘NWC’ 时形状为 [batch, in_width, in_channel]

  • kernel (tvm.te.Tensor) – 3-D 卷积核,当 kernel_layout == ‘OIW’ 时形状为 [num_filter, in_channel, filter_size],当 kernel_layout == ‘WIO’ 时形状为 [filter_size, in_channel, num_filter]

  • strides (inttuple) – 沿宽度的空间步幅

  • padding (intstr) – 填充大小,或 [‘VALID’, ‘SAME’]

  • dilation (inttuple) – 如果应使用空洞卷积,则为空洞率。

  • data_layout (str) – 输入数据如何布局,必须是 [‘NCW’, ‘NWC’] 之一

  • kernel_layout (Optiona[str]) – 卷积核的布局。如果未指定,则使用默认布局。如果 data_layout == “NCW”,则为 “OIW”;如果 data_layout == “NWC”,则为 “WIO”。

  • out_dtype (str) – 输出数据类型。如果为 None,则输出与输入类型相同。

tvm.topi.nn.conv1d_ncw(data, kernel, strides=1, padding='VALID', dilation=1, out_dtype=None)

NCW 布局中的 1D 卷积。有关参数的详细信息,请参阅 conv()

tvm.topi.nn.conv1d_nwc(data, kernel, strides=1, padding='VALID', dilation=1, out_dtype=None)

NWC 布局中的 1D 卷积。有关参数的详细信息,请参阅 conv()

tvm.topi.nn.conv1d_transpose_ncw(data, kernel, stride, padding, out_dtype, output_padding)

转置 1D 卷积 ncw 前向运算符。

参数:
  • data (tvm.te.Tensor) – 3-D 张量,形状为 [batch, in_channel, in_width]

  • kernel (tvm.te.Tensor) – 3-D 张量,形状为 [in_channel, num_filter, filter_width]

  • stride (ints) – 沿宽度的空间步幅

  • padding (intstr) – 填充大小,或 [‘VALID’, ‘SAME’]

  • out_dtype (str) – 输出数据类型。用于混合精度。

  • output_padding (ints) – 用于在存在多个可能的形状时恢复实际输出形状。必须小于步幅。

返回:

output – 3-D 张量,形状为 [batch, out_channel, out_width]

返回类型:

tvm.te.Tensor

tvm.topi.nn.conv2d(input, filter, strides, padding, dilation, data_layout='NCHW', kernel_layout='', out_dtype=None)

Conv2D 运算符。

参数:
  • input (tvm.te.Tensor) – 4-D 张量,在 data_layout 布局中形状为 [batch, in_channel, in_height, in_width]

  • filter (tvm.te.Tensor) – 4-D 张量,在 kernel_layout 布局中形状为 [num_filter, in_channel, filter_height, filter_width]

  • strides (int两个整数的列表/元组) – 步幅大小,或 [stride_height, stride_width]

  • padding (int2 个或 4 个整数的列表/元组) – 填充大小,或 2 个整数的 [pad_height, pad_width],或 4 个整数的 [pad_top, pad_left, pad_bottom, pad_right]

  • dilation (int两个整数的列表/元组) – 空洞大小,或 [dilation_height, dilation_width]

  • data_layout (str) – 数据布局

  • kernel_layout (Optional[str]) – 卷积核布局。如果未指定,则使用从 data_layout 推断的默认布局。如果 data_layout == “NCHW”,则为 “OIHW”;如果 data_layout == “NHWC”,则为 “HWIO”。

返回:

output – 4 维,形状为 [batch, out_channel, out_height, out_width]

返回类型:

tvm.te.Tensor

tvm.topi.nn.conv2d_NCHWc(data, kernel, stride, padding, dilation, layout, out_layout, out_dtype='float32')

用于 nChw[x]c 布局的 Conv2D 运算符。

参数:
  • data (tvm.te.Tensor) – 5-D 张量,形状为 [batch, in_channel_chunk, in_height, in_width, in_channel_block]

  • kernel (tvm.te.Tensor) – 6-D 张量,形状为 [num_filter_chunk, in_channel_chunk, filter_height, filter_width, in_channel_block, num_filter_block]

  • stride (int or a list/tuple of two ints) – 步幅大小,或 [stride_height, stride_width]

  • padding (int2 个或 4 个整数的列表/元组) – 填充大小,或 2 个整数的 [pad_height, pad_width],或 4 个整数的 [pad_top, pad_left, pad_bottom, pad_right]

  • dilation (int两个整数的列表/元组) – 空洞大小,或 [dilation_height, dilation_width]

  • layout (str) – 输入数据布局

  • out_layout (str) – 输出数据布局

  • out_dtype (str) – 输出数据类型

返回:

output – 5-D 张量,形状为 [batch, out_channel_chunk, out_height, out_width, out_channel_block]

返回类型:

tvm.te.Tensor

tvm.topi.nn.conv2d_NCHWc_int8(data, kernel, stride, padding, dilation, layout, out_layout, out_dtype='int32', n_elems=4)

用于 nChw[x]c 布局的 Conv2D 运算符。

参数:
  • data (tvm.te.Tensor) – 5-D 张量,形状为 [batch, in_channel_chunk, in_height, in_width, in_channel_block]

  • kernel (tvm.te.Tensor) – 7-D 张量,形状为 [num_filter_chunk, in_channel_chunk, filter_height, filter_width, in_channel_block/4, num_filter_block, 4]

  • stride (int or a list/tuple of two ints) – 步幅大小,或 [stride_height, stride_width]

  • padding (int2 个或 4 个整数的列表/元组) – 填充大小,或 2 个整数的 [pad_height, pad_width],或 4 个整数的 [pad_top, pad_left, pad_bottom, pad_right]

  • dilation (int两个整数的列表/元组) – 空洞大小,或 [dilation_height, dilation_width]

  • layout (str) – 输入数据布局

  • out_layout (str) – 输出数据布局

  • out_dtype (str) – 输出数据类型

  • n_elems (int) – 累积的 int8 元素的数量

返回:

output – 5-D 张量,形状为 [batch, out_channel_chunk, out_height, out_width, out_channel_block]

返回类型:

tvm.te.Tensor

tvm.topi.nn.conv2d_hwcn(Input, Filter, stride, padding, dilation, out_dtype=None)

HWCN 布局中的卷积运算符。

参数:
  • Input (tvm.te.Tensor) – 4-D 张量,形状为 [in_height, in_width, in_channel, batch]

  • Filter (tvm.te.Tensor) – 4-D 张量,形状为 [filter_height, filter_width, in_channel, num_filter]

  • stride (int两个整数的列表/元组) – 步幅大小,或 [stride_height, stride_width]

  • padding (int2 个或 4 个整数的列表/元组) – 填充大小,或 2 个整数的 [pad_height, pad_width],或 4 个整数的 [pad_top, pad_left, pad_bottom, pad_right]

  • dilation (int两个整数的列表/元组) – 空洞大小,或 [dilation_height, dilation_width]

返回:

output – 4-D 张量,形状为 [out_height, out_width, out_channel, batch]

返回类型:

tvm.te.Tensor

tvm.topi.nn.conv2d_nchw(Input, Filter, stride, padding, dilation, out_dtype=None)

NCHW 布局中的卷积运算符。

参数:
  • Input (tvm.te.Tensor) – 4-D 张量,形状为 [batch, in_channel, in_height, in_width]

  • Filter (tvm.te.Tensor) – 4-D 张量,形状为 [num_filter, in_channel, filter_height, filter_width]

  • stride (int两个整数的列表/元组) – 步幅大小,或 [stride_height, stride_width]

  • padding (int2 个或 4 个整数的列表/元组) – 填充大小,或 2 个整数的 [pad_height, pad_width],或 4 个整数的 [pad_top, pad_left, pad_bottom, pad_right]

  • dilation (int两个整数的列表/元组) – 空洞大小,或 [dilation_height, dilation_width]

返回:

Output – 4-D 张量,形状为 [batch, out_channel, out_height, out_width]

返回类型:

tvm.te.Tensor

tvm.topi.nn.conv2d_nhwc(Input, Filter, stride, padding, dilation, out_dtype='float32', auto_scheduler_rewritten_layout='', meta_schedule_original_shape=None)

NHWC 布局中的卷积运算符。

参数:
  • Input (tvm.te.Tensor) – 4-D 张量,形状为 [batch, in_height, in_width, in_channel]

  • Filter (tvm.te.Tensor) – 4-D 张量,形状为 [filter_height, filter_width, in_channel, num_filter]

  • stride (int两个整数的列表/元组) – 步幅大小,或 [stride_height, stride_width]

  • padding (int2 个或 4 个整数的列表/元组) – 填充大小,或 2 个整数的 [pad_height, pad_width],或 4 个整数的 [pad_top, pad_left, pad_bottom, pad_right]

  • dilation (int两个整数的列表/元组) – 空洞大小,或 [dilation_height, dilation_width]

  • out_dtype (str = "float32",) – 输出张量的类型

  • auto_scheduler_rewritten_layout (str = "") – 自动调度器的布局重写过程之后的布局。

  • meta_schedule_original_shape (Optional[List[PrimExpr]] = None) – 输入张量的原始形状。

返回:

output – 4 维,形状为 [batch, out_height, out_width, out_channel]

返回类型:

tvm.te.Tensor

tvm.topi.nn.conv2d_transpose_nchw(Input, Filter, strides, padding, out_dtype, output_padding)

转置 2D 卷积 nchw 前向运算符。

参数:
  • Input (tvm.te.Tensor) – 4-D 张量,形状为 [batch, in_channel, in_height, in_width]

  • Filter (tvm.te.Tensor) – 4-D 张量,形状为 [in_channel, num_filter, filter_height, filter_width]

  • strides (tuple of two ints) – 沿高度和宽度的空间步幅

  • padding (intstr) – 填充大小,或 [‘VALID’, ‘SAME’]

  • out_dtype (str) – 输出数据类型。用于混合精度。

  • output_padding (tuple of ints) – 用于获取正确的梯度输出形状

返回:

Output – 4-D 张量,形状为 [batch, out_channel, out_height, out_width]

返回类型:

tvm.te.Tensor

tvm.topi.nn.conv2d_transpose_nchw_preprocess(data, kernel, strides, padding, out_dtype, output_padding)

预处理数据和内核,使 conv2d_transpose 的计算模式与 conv2d 相同

tvm.topi.nn.conv2d_winograd_nchw(data, weight, strides, padding, dilation, out_dtype, pre_computed=False, auto_scheduler_rewritten_layout='', meta_schedule_original_shape=None)

NCHW 布局中的 Conv2D Winograd 卷积。这是一个清晰的版本,供自动调度器在 CPU 和 GPU 上使用。

参数:
  • data (tvm.te.Tensor) – 4 维,形状为 [batch, in_channel, in_height, in_width]

  • weight (tvm.te.Tensor) – 4-D 张量,形状为 [filter_height, filter_width, in_channel, num_filter]

  • strides (int两个整数的列表/元组) – 步幅大小,或 [stride_height, stride_width]

  • padding (int两个整数的列表/元组) – 填充大小,或 [pad_height, pad_width]

  • dilation (int两个整数的列表/元组) – 空洞大小,或 [dilation_height, dilation_width]

  • out_dtype (str, 可选) – 指定输出数据类型。

  • pre_computed (bool) – 卷积核是否为预计算的

  • auto_scheduler_rewritten_layout (str = "") – 自动调度器的布局重写过程之后的布局。

  • meta_schedule_original_shape (Optional[List[PrimExpr]] = None) – 输入张量的原始形状。

返回:

output – 4 维,形状为 [batch, out_height, out_width, out_channel]

返回类型:

tvm.te.Tensor

tvm.topi.nn.conv2d_winograd_nchw_without_weight_transform(data, weight, strides, padding, dilation, out_dtype, auto_scheduler_rewritten_layout='', meta_schedule_original_shape=None)

NCHW 布局中不带布局变换的 Conv2D Winograd 卷积。这是一个清晰的版本,供元调度在 CPU 和 GPU 上使用。

参数:
  • data (tvm.te.Tensor) – 4 维,形状为 [batch, in_height, in_width, in_channel]

  • weight (tvm.te.Tensor) – 4-D 张量,形状为 [filter_height, filter_width, in_channel, num_filter]

  • strides (int两个整数的列表/元组) – 步幅大小,或 [stride_height, stride_width]

  • padding (int两个整数的列表/元组) – 填充大小,或 [pad_height, pad_width]

  • dilation (int两个整数的列表/元组) – 空洞大小,或 [dilation_height, dilation_width]

  • out_dtype (str, 可选) – 指定输出数据类型。

  • auto_scheduler_rewritten_layout (str = "") – 自动调度器的布局重写过程之后的布局。

  • meta_schedule_original_shape (Optional[List[PrimExpr]] = None) – 输入张量的原始形状。

返回:

output – 4 维,形状为 [batch, out_height, out_width, out_channel]

返回类型:

tvm.te.Tensor

tvm.topi.nn.conv2d_winograd_nhwc(data, weight, strides, padding, dilation, out_dtype, pre_computed=False, auto_scheduler_rewritten_layout='', meta_schedule_original_shape=None)

NHWC 布局中的 Conv2D Winograd 卷积。这是一个清晰的版本,供自动调度器在 CPU 和 GPU 上使用。

参数:
  • data (tvm.te.Tensor) – 4 维,形状为 [batch, in_height, in_width, in_channel]

  • weight (tvm.te.Tensor) – 4-D 张量,形状为 [filter_height, filter_width, in_channel, num_filter]

  • strides (int两个整数的列表/元组) – 步幅大小,或 [stride_height, stride_width]

  • padding (int两个整数的列表/元组) – 填充大小,或 [pad_height, pad_width]

  • dilation (int两个整数的列表/元组) – 空洞大小,或 [dilation_height, dilation_width]

  • out_dtype (str, 可选) – 指定输出数据类型。

  • pre_computed (bool) – 卷积核是否为预计算的

  • auto_scheduler_rewritten_layout (str = "") – 自动调度器的布局重写过程之后的布局。

  • meta_schedule_original_shape (Optional[List[PrimExpr]] = None) – 输入张量的原始形状。

返回:

output – 4 维,形状为 [batch, out_height, out_width, out_channel]

返回类型:

tvm.te.Tensor

tvm.topi.nn.conv2d_winograd_nhwc_without_weight_transform(data, weight, strides, padding, dilation, out_dtype, auto_scheduler_rewritten_layout='', meta_schedule_original_shape=None)

NHWC 布局中不带布局变换的 Conv2D Winograd 卷积。这是一个清晰的版本,供自动调度器在 CPU 和 GPU 上使用。

参数:
  • data (tvm.te.Tensor) – 4 维,形状为 [batch, in_height, in_width, in_channel]

  • weight (tvm.te.Tensor) – 4-D 张量,形状为 [filter_height, filter_width, in_channel, num_filter]

  • strides (int两个整数的列表/元组) – 步幅大小,或 [stride_height, stride_width]

  • padding (int两个整数的列表/元组) – 填充大小,或 [pad_height, pad_width]

  • dilation (int两个整数的列表/元组) – 空洞大小,或 [dilation_height, dilation_width]

  • out_dtype (str, 可选) – 指定输出数据类型。

  • auto_scheduler_rewritten_layout (str = "") – 自动调度器的布局重写过程之后的布局。

  • meta_schedule_original_shape (Optional[List[PrimExpr]] = None) – 输入张量的原始形状。

返回:

output – 4 维,形状为 [batch, out_height, out_width, out_channel]

返回类型:

tvm.te.Tensor

tvm.topi.nn.conv2d_winograd_weight_transform(kernel, tile_size)

winograd 的权重变换

参数:
  • kernel (Tensor) – 原始卷积核张量,布局为 “NCHW”。

  • tile_size (int) – Winograd 变换的分块大小。例如,F(2x2, 3x3) 为 2,F(4x4, 3x3) 为 4

返回:

output – 4-D 张量,形状为 [alpha, alpha, CO, CI]

返回类型:

tvm.te.Tensor

tvm.topi.nn.conv3d_ncdhw(Input, Filter, stride, padding, dilation, groups, out_dtype=None)

NCDHW 布局中的 Conv3D 运算符。

参数:
  • Input (tvm.te.Tensor) – 5 维张量,形状为 [batch, in_channel, in_depth, in_height, in_width]

  • Filter (tvm.te.Tensor) – 5 维张量,形状为 [num_filter, in_channel, filter_depth, filter_height, filter_width]

  • stride (int一个包含三个整数的列表/元组) – 步长大小,或 [strid_depth, stride_height, stride_width]

  • padding (intstr) – 填充大小,或 [‘VALID’, ‘SAME’]

  • dilation (int一个包含三个整数的列表/元组) – 空洞大小,或 [dilation_depth, dilation_height, dilation_width]

  • groups (int) – 分组数。

返回:

Output – 5 维张量,形状为 [batch, out_channel, out_depth, out_height, out_width]

返回类型:

tvm.te.Tensor

tvm.topi.nn.conv3d_ndhwc(Input, Filter, stride, padding, dilation, groups, out_dtype='float32', auto_scheduler_rewritten_layout='', meta_schedule_origin_shape=None)

NDHWC 布局中的卷积运算符。

参数:
  • Input (tvm.te.Tensor) – 5 维张量,形状为 [batch, in_depth, in_height, in_width, in_channel]

  • Filter (tvm.te.Tensor) – 5 维张量,形状为 [filter_depth, filter_height, filter_width, in_channel, num_filter]

  • stride (int一个包含三个整数的列表/元组) – 步长大小,或 [stride_depth, stride_height, stride_width]

  • padding (intstr) – 填充大小,或 [‘VALID’, ‘SAME’]

  • dilation (int一个包含三个整数的列表/元组) – 空洞大小,或 [dilation_depth, dilation_height, dilation_width]

  • groups (int) – 分组数。

  • out_dtype (str = "float32",) – 输出张量的类型

  • auto_scheduler_rewritten_layout (str = "") – 自动调度器的布局重写过程之后的布局。

  • meta_schedule_origin_shape (Optional[List[PrimExpr]] = None) – 输入张量的原始形状。

返回:

Output – 5 维张量,形状为 [batch, out_depth, out_height, out_width, out_channel]

返回类型:

tvm.te.Tensor

tvm.topi.nn.conv3d_transpose_ncdhw(Input, Filter, strides, padding, out_dtype, output_padding)

转置 3D 卷积 ncdhw 前向运算符。

参数:
  • Input (tvm.te.Tensor) – 5 维张量,形状为 [batch, in_channel, in_depth, in_height, in_width]

  • Filter (tvm.te.Tensor) – 5 维张量,形状为 [in_channel, num_filter, filter_depth, filter_height, filter_width]

  • strides (int一个包含三个整数的列表/元组) – 沿深度、高度和宽度的空间步长

  • padding (intstr) – 填充大小,或 [‘VALID’, ‘SAME’]

  • out_dtype (str) – 输出数据类型。用于混合精度。

  • output_padding (tuple of ints) – 用于获取正确的梯度输出形状

返回:

Output – 5 维张量,形状为 [batch, out_channel, out_depth, out_height, out_width]

返回类型:

tvm.te.Tensor

tvm.topi.nn.conv3d_transpose_ncdhw_preprocess(data, kernel, strides, padding, out_dtype, output_padding)

预处理数据和内核,使 conv3d_transpose 的计算模式与 conv3d 相同

tvm.topi.nn.conv3d_winograd_weight_transform(kernel, tile_size)

3D winograd 的权重变换

参数:
  • kernel (Tensor) – 原始核张量,布局为 “NCDHW”。

  • tile_size (int) – Winograd 变换的分块大小。例如,F(2x2, 3x3) 为 2,F(4x4, 3x3) 为 4

返回:

output – 5 维张量,形状为 [alpha, alpha, alpha, CO, CI]

返回类型:

tvm.te.Tensor

tvm.topi.nn.correlation_nchw(data1, data2, kernel_size, max_displacement, stride1, stride2, padding, is_multiply)

NCHW 布局中的相关运算符。

参数:
  • data1 (tvm.te.Tensor) – 4 维张量,形状为 [batch, channel, height, width]

  • data2 (tvm.te.Tensor) – 4 维张量,形状为 [batch, channel, height, width]

  • kernel_size (int) – 相关性的核大小,必须是奇数

  • max_displacement (int) – 相关性的最大位移

  • stride1 (int) – data1 的步长

  • stride2 (int) – 在以 data1 为中心的邻域内,data2 的步长

  • padding (int一个包含 2 个或 4 个整数的列表/元组) – 填充大小,对于 2 个整数为 [pad_height, pad_width],对于 4 个整数为 [pad_top, pad_left, pad_bottom, pad_right]

  • is_multiply (bool) – 运算类型是乘法还是减法

返回:

Output – 4-D 张量,形状为 [batch, out_channel, out_height, out_width]

返回类型:

tvm.te.Tensor

tvm.topi.nn.declaration_conv2d_transpose_impl(data, kernel, strides, padding, out_dtype, output_padding)

conv2d 转置的实现

tvm.topi.nn.declaration_conv3d_transpose_impl(data, kernel, strides, padding, out_dtype, output_padding)

conv3d 转置的实现

tvm.topi.nn.deformable_conv2d_nchw(data, offset, kernel, strides, padding, dilation, deformable_groups, groups, out_dtype)

NCHW 布局中的可变形 conv2D 运算符。

可变形卷积操作的描述见 https://arxiv.org/abs/1703.06211

参数:
  • data (tvm.te.Tensor) – 4 维,形状为 [batch, in_channel, in_height, in_width]

  • offset (tvm.te.Tensor) – 4 维张量,形状为 [batch, deformable_groups * filter_height * filter_width * 2, out_height, out_width]。

  • kernel (tvm.te.Tensor) – 4 维,形状为 [num_filter, in_channel, filter_height, filter_width]

  • strides (int两个整数的列表/元组) – 步幅大小,或 [stride_height, stride_width]

  • padding (int两个整数的列表/元组) – 填充大小,或 [pad_height, pad_width]

  • dilation (int两个整数的列表/元组) – 空洞大小,或 [dilation_height, dilation_width]

  • deformable_groups (int) – 可变形分组的数量

  • groups (int) – 分组数

返回:

output – 4 维,形状为 [batch, out_channel, out_height, out_width]

返回类型:

tvm.te.Tensor

tvm.topi.nn.deformable_conv2d_nhwc(data, offset, kernel, strides, padding, dilation, deformable_groups, groups, out_dtype)

NHWC 布局中的可变形 conv2D 运算符。

可变形卷积操作的描述见 https://arxiv.org/abs/1703.06211

参数:
  • data (tvm.te.Tensor) – 4 维,形状为 [batch, in_height, in_width, in_channel]

  • offset (tvm.te.Tensor) –

    4 维张量,形状为 [batch, out_height, out_width,

    deformable_groups * filter_height * filter_width * 2]。

  • kernel (tvm.te.Tensor) – 4 维,形状为 [filter_height, filter_width, in_channel, num_filter]

  • strides (int两个整数的列表/元组) – 步幅大小,或 [stride_height, stride_width]

  • padding (int两个整数的列表/元组) – 填充大小,或 [pad_height, pad_width]

  • dilation (int两个整数的列表/元组) – 空洞大小,或 [dilation_height, dilation_width]

  • deformable_groups (int) – 可变形分组的数量

  • groups (int) – 分组数

返回:

output – 4 维,形状为 [batch, out_height, out_width, out_channel]

返回类型:

tvm.te.Tensor

tvm.topi.nn.dense(data, weight, bias=None, out_dtype=None, auto_scheduler_rewritten_layout='', meta_schedule_original_shape=None)

topi 中 dense 的默认实现。这是 matmul_nt 运算符的别名,用于非转置格式的数据张量和转置格式的权重张量。

参数:
  • data (tvm.te.Tensor) – 2 维,形状为 [batch, in_dim]

  • weight (tvm.te.Tensor) – 2 维张量,形状为 [out_dim, in_dim]

  • bias (Optional[tvm.te.Tensor]) – 1 维张量,形状为 [out_dim]

  • out_dtype (Optional[str]) – 输出类型。用于混合精度。

  • auto_scheduler_rewritten_layout (str = "") – 自动调度器的布局重写过程之后的布局。

  • meta_schedule_original_shape (Optional[List[PrimExpr]] = None) – 输入张量的原始形状。

返回:

output – 2 维,形状为 [batch, out_dim]

返回类型:

tvm.te.Tensor

tvm.topi.nn.dense_pack(data, weight, bias=None, out_dtype=None)

topi 中 dense_pack 的默认实现。

参数:
  • data (tvm.te.Tensor) – 2 维,形状为 [batch, in_dim]

  • weight (tvm.te.Tensor) – 2 维张量,形状为 [out_dim, in_dim]

  • bias (Optional[tvm.te.Tensor]) – 1 维张量,形状为 [out_dim]

  • out_dtype (Optional[str]) – 输出类型。用于混合精度。

返回:

output – 2 维,形状为 [batch, out_dim]

返回类型:

tvm.te.Tensor

tvm.topi.nn.depth_to_space(data, block_size, layout='NCHW', mode='DCR')

对数据执行深度到空间转换

参数:
  • data (tvm.te.Tensor) – 4 维张量,布局为 NCHW 或 NHWC。

  • block_size (int) – 从通道维度组合的块大小。

  • layout (string) – NCHW 或 NHWC,表示数据布局。

  • mode (string) – DCR 或 CDR,表示应如何访问通道。在 DCR 中,通道以 Tensorflow 风格交织,而在 CDR 中,通道像 Pytorch 中一样按顺序访问。

返回:

output – 输出形状为 [N, C / block_size**2, H * block_size, W * block_size]

返回类型:

tvm.te.Tensor

tvm.topi.nn.depthwise_conv2d_NCHWc(Input, Filter, stride, padding, dilation, layout, out_layout, out_dtype=None)

深度可分离卷积 NCHW[x]c 前向运算符。

参数:
  • Input (tvm.te.Tensor) – 5 维张量,形状为 [batch, in_channel_chunk, in_height, in_width, in_channel_block]

  • Filter (tvm.te.Tensor) – 6 维张量,形状为 [out_channel_chunk, 1, filter_height, filter_width, 1, out_channel_block]。在 NCHWc 深度卷积中,我们将核的 in_channel 和 channel_multiplier 分组在一起,然后进行平铺。

  • stride (tuple of two ints) – 沿高度和宽度的空间步长

  • padding (intstr) – 填充大小,或 [‘VALID’, ‘SAME’]

  • dilation (int两个整数的列表/元组) – 空洞大小,或 [dilation_height, dilation_width]

  • layout (str) – 输入数据布局

  • out_layout (str) – 输出数据布局

  • out_dtype (str, optional) – 输出数据类型

返回:

Output – 5 维张量,形状为 [batch, out_channel_chunk, out_height, out_width, out_channel_block]

返回类型:

tvm.te.Tensor

tvm.topi.nn.depthwise_conv2d_backward_input_nhwc(Filter, Out_grad, oshape, ishape, stride, padding)

深度可分离卷积 nhwc 反向传播输入运算符。

参数:
  • Filter (tvm.te.Tensor) – 4 维张量,形状为 [filter_height, filter_width, in_channel, channel_multiplier]

  • Out_grad (tvm.te.Tensor) – 4 维张量,形状为 [batch, out_height, out_width, out_channel]

  • stride (tuple of two ints) – 沿高度和宽度的空间步长

  • padding (intstr) – 填充大小,或 [‘VALID’, ‘SAME’]

返回:

Output – 4 维张量,形状为 [batch, in_height, in_width, in_channel]

返回类型:

tvm.te.Tensor

tvm.topi.nn.depthwise_conv2d_backward_weight_nhwc(Input, Out_grad, oshape, fshape, stride, padding)

深度可分离卷积 nhwc 反向传播权重运算符。

参数:
  • Input (tvm.te.Tensor) – 4-D 张量,形状为 [batch, in_height, in_width, in_channel]

  • Out_grad (tvm.te.Tensor) – 4 维张量,形状为 [batch, out_height, out_width, out_channel]

  • stride (tuple of two ints) – 沿高度和宽度的空间步长

  • padding (intstr) – 填充大小,或 [‘VALID’, ‘SAME’]

返回:

Output – 4 维张量,形状为 [filter_height, filter_width, in_channel, channel_multiplier]

返回类型:

tvm.te.Tensor

tvm.topi.nn.depthwise_conv2d_nchw(Input, Filter, stride, padding, dilation, out_dtype=None)

深度可分离卷积 nchw 前向运算符。

参数:
  • Input (tvm.te.Tensor) – 4-D 张量,形状为 [batch, in_channel, in_height, in_width]

  • Filter (tvm.te.Tensor) – 4 维张量,形状为 [in_channel, channel_multiplier, filter_height, filter_width]

  • stride (int一个包含两个整数的列表/元组) – 空间步长,或 (stride_height, stride_width)。

  • padding (intstr) – 填充大小,或 [‘VALID’, ‘SAME’]

  • dilation (int两个整数的列表/元组) – 空洞大小,或 [dilation_height, dilation_width]

  • out_dtype (str, optional) – 输出数据类型

返回:

Output – 4-D 张量,形状为 [batch, out_channel, out_height, out_width]

返回类型:

tvm.te.Tensor

tvm.topi.nn.depthwise_conv2d_nhwc(Input, Filter, stride, padding, dilation, kernel_layout='HWOI', out_dtype=None)

深度可分离卷积 nhwc 前向运算符。

参数:
  • Input (tvm.te.Tensor) – 4-D 张量,形状为 [batch, in_height, in_width, in_channel]

  • Filter (tvm.te.Tensor) – 4 维张量,形状为 [filter_height, filter_width, in_channel, channel_multiplier]

  • stride (tuple of two ints) – 沿高度和宽度的空间步长

  • padding (intstr) – 填充大小,或 [‘VALID’, ‘SAME’]

  • dilation (int两个整数的列表/元组) – 空洞大小,或 [dilation_height, dilation_width]

  • out_dtype (str, optional) – 输出数据类型

返回:

Output – 4 维张量,形状为 [batch, out_height, out_width, out_channel]

返回类型:

tvm.te.Tensor

tvm.topi.nn.dilate(data, strides, dilation_value=0.0, name='DilatedInput')

使用给定的膨胀值(默认为 0)膨胀数据。

参数:
  • data (tvm.te.Tensor) – n 维张量,可以是任何布局。

  • strides (list / tuple of n ints) – 每个维度上的空洞步长,1 表示无空洞。

  • dilation_value (int/float, optional) – 用于空洞输入的值。

  • name (str, optional) – 生成的运算符的名称前缀

返回:

Output – n 维张量,与 data 相同的布局。

返回类型:

tvm.te.Tensor

tvm.topi.nn.equal_const_int(expr, value)

返回 expr 是否等于 value。

参数:

expr (tvm.Expr) – 输入表达式。

返回:

equal – 是否相等。

返回类型:

bool

tvm.topi.nn.fast_softmax(x, axis=-1)

对数据执行 softmax 激活。使用近似值来计算指数,以提高速度。

参数:
返回:

output – 输出形状与输入相同

返回类型:

tvm.te.Tensor

tvm.topi.nn.fifo_buffer(data, buffer, axis)

FIFO 缓冲区,用于在具有滑动窗口输入的 CNN 中实现计算重用

计算等价于

concat(buffer, data, axis=axis)
.slice_axis(axis=axis,
            begin=data.shape[axis],
            end=data.shape[axis]+buffer.shape[axis])

适用于

  • 在滑动窗口输入上操作的卷积运算中,显式编码计算的重用

  • 实现 FIFO 队列以缓存中间结果,例如 Fast WaveNet 中的那样。

参数:
返回:

result – 缓冲区的更新值

返回类型:

tvm.te.Tensor

tvm.topi.nn.flatten(data)

通过折叠更高维度,将输入数组展平为 2-D 数组。

参数:

data (tvm.te.Tensor) – 输入数组。

返回:

output – 2 维数组,具有折叠的更高维度。

返回类型:

tvm.te.Tensor

tvm.topi.nn.get_const_int(expr)

验证 expr 是否为整数并获取常量值。

参数:

expr (tvm.Expr or int) – 输入表达式。

返回:

out_value – 输出值。

返回类型:

int

tvm.topi.nn.get_const_tuple(in_tuple)

验证输入元组是否为 IntImm 或 Var,返回 int 或 Var 元组。

参数:

in_tuple (tuple of Expr) – 输入。

返回:

out_tuple – 输出。

返回类型:

tuple of int

tvm.topi.nn.get_pad_tuple(padding, kernel)

获取填充选项的通用代码

参数:
  • padding (intstr) – 填充大小,或 [‘VALID’, ‘SAME’]

  • kernel (tuple of int) – 卷积核大小

返回:

  • pad_top (int) – 顶部的填充大小

  • pad_left (int) – 左侧的填充大小

  • pad_down (int) – 底部的填充大小。

  • pad_right (int) – 右侧的填充大小。

tvm.topi.nn.get_pad_tuple1d(padding, kernel)

获取填充选项的通用代码

参数:
  • padding (intstr) – 填充大小,或 [‘VALID’, ‘SAME’]

  • kernel (tuple of int) – 卷积核大小

返回:

  • pad_left (int) – 左侧的填充大小

  • pad_right (int) – 右侧的填充大小。

tvm.topi.nn.get_pad_tuple3d(padding, kernel)

获取填充选项的通用代码

参数:
  • padding (intstr) – 填充大小,或 [‘VALID’, ‘SAME’]

  • kernel (tuple of int) – 卷积核大小

返回:

  • pad_front (int) – 前面的填充大小。

  • pad_top (int) – 顶部的填充大小

  • pad_left (int) – 左侧的填充大小

  • pad_back (int) – 后面的填充大小。

  • pad_down (int) – 底部的填充大小。

  • pad_right (int) – 右侧的填充大小。

tvm.topi.nn.get_pad_tuple_generic(padding, kernel)

获取填充选项的通用代码

参数:
  • padding (intstr) – 填充大小,或 [‘VALID’, ‘SAME’]

  • kernel (tuple of int) – 卷积核大小

返回:

  • pad_top (int) – 顶部的填充大小

  • pad_down (int) – 底部的填充大小。

  • pad_left (int) – 左侧的填充大小

  • pad_right (int) – 右侧的填充大小。

tvm.topi.nn.global_pool(data, pool_type, layout='NCHW')
对数据的高度和宽度维度执行全局池化。

它根据布局字符串决定高度和宽度维度,其中 ‘W’ 和 ‘H’ 分别表示宽度和高度。宽度和高度维度不能拆分。例如,NCHW、NCHW16c 等对于池化有效,而 NCHW16w、NCHW16h 则无效。有关布局字符串约定的更多信息,请参见 layout 参数。

参数:
  • data (tvm.te.Tensor) – n 维,形状为布局

  • pool_type (str) – 池化类型,‘max’ 或 ‘avg’

  • layout (str) – 输入数据的布局。布局应由大写字母、小写字母和数字组成,其中大写字母表示维度,相应的小写字母和因子大小表示拆分维度。例如,NCHW16c 可以描述一个 5 维张量 [batch_size, channel, height, width, channel_block],其中 channel_block=16 是维度通道的拆分。

返回:

output – 与输入布局相同的 n 维张量,高度和宽度维度大小为 1。例如,对于 NCHW,输出形状将为 [batch, channel, 1, 1]

返回类型:

tvm.te.Tensor

tvm.topi.nn.group_conv1d_ncw(data, kernel, strides=1, padding='VALID', dilation=1, groups=1, out_dtype=None)

NCW 布局的 1D 卷积前向运算符。

参数:
  • data (tvm.te.Tensor) – 3-D 张量,形状为 [batch, in_channel, in_width]

  • kernel (tvm.te.Tensor) – 3 维张量,形状为 [num_filter, in_channel, filter_size]

  • strides (inttuple) – 沿宽度的空间步幅

  • padding (int, tuple, or str) – 填充大小可以是整数(表示等量填充)、(左,右) 元组或 [‘VALID’, ‘SAME’] 字符串。

  • dilation (inttuple) – 如果应使用空洞卷积,则为空洞率。

  • groups (int) – 组数

  • out_dtype (str) – 输出数据类型。如果为 None,则输出与输入类型相同。

tvm.topi.nn.group_conv1d_nwc(data, kernel, strides=1, padding='VALID', dilation=1, groups=1, out_dtype=None)

NWC 布局的 1D 卷积前向运算符。

参数:
  • data (tvm.te.Tensor) – 3 维,形状为 [batch, in_width, in_channel]

  • kernel (tvm.te.Tensor) – 3 维,形状为 [filter_size, in_channel, num_filter]

  • strides (inttuple) – 沿宽度的空间步幅

  • padding (int, tuple, or str) – 填充大小可以是整数(表示等量填充)、(左,右) 元组或 [‘VALID’, ‘SAME’] 字符串。

  • dilation (inttuple) – 如果应使用空洞卷积,则为空洞率。

  • groups (int) – 组数

  • out_dtype (str) – 输出数据类型。如果为 None,则输出与输入类型相同。

tvm.topi.nn.group_conv1d_transpose_ncw(data, kernel, stride, padding, out_dtype, output_padding, groups)

转置 1D 分组卷积 ncw 前向运算符。

参数:
  • data (tvm.te.Tensor) – 3-D 张量,形状为 [batch, in_channel, in_width]

  • kernel (tvm.te.Tensor) – 3-D 张量,形状为 [in_channel, num_filter, filter_width]

  • stride (ints) – 沿宽度的空间步幅

  • padding (intstr) – 填充大小,或 [‘VALID’, ‘SAME’]

  • out_dtype (str) – 输出数据类型。用于混合精度。

  • output_padding (ints) –

    用于在存在多个可能形状的情况下恢复实际输出形状。

    必须小于步幅 (stride)。

    groupsint

    组数

返回:

output – 3-D 张量,形状为 [batch, out_channel, out_width]

返回类型:

tvm.te.Tensor

tvm.topi.nn.group_conv2d_nchw(Input, Filter, stride, padding, dilation, groups, out_dtype=None)

NCHW 布局中的分组卷积运算符。

参数:
  • Input (tvm.te.Tensor) – 4-D 张量,形状为 [batch, in_channel, in_height, in_width]

  • Filter (tvm.te.Tensor) – 4 维,形状为 [num_filter, in_channel // groups, filter_height, filter_width]

  • stride (int两个整数的列表/元组) – 步幅大小,或 [stride_height, stride_width]

  • padding (int2 个或 4 个整数的列表/元组) – 填充大小,或 2 个整数的 [pad_height, pad_width],或 4 个整数的 [pad_top, pad_left, pad_bottom, pad_right]

  • dilation (int两个整数的列表/元组) – 空洞大小,或 [dilation_height, dilation_width]

  • groups (int) – 分组数

  • out_dtype (str) – 输出类型。用于混合精度。

返回:

Output – 4-D 张量,形状为 [batch, out_channel, out_height, out_width]

返回类型:

tvm.te.Tensor

tvm.topi.nn.group_conv2d_nhwc(Input, Filter, stride, padding, dilation, groups, out_dtype=None)

NHWC 布局中的分组卷积运算符。

参数:
  • Input (tvm.te.Tensor) – 4 维,形状为 [batch, in_height, in_width, in_channel, …]

  • Filter (tvm.te.Tensor) – 4 维,形状为 [filter_height, filter_width, in_channel // groups, num_filter]

  • stride (int两个整数的列表/元组) – 步幅大小,或 [stride_height, stride_width]

  • padding (int2 个或 4 个整数的列表/元组) – 填充大小,或 2 个整数的 [pad_height, pad_width],或 4 个整数的 [pad_top, pad_left, pad_bottom, pad_right]

  • dilation (int两个整数的列表/元组) – 空洞大小,或 [dilation_height, dilation_width]

  • groups (int) – 分组数

  • out_dtype (str) – 输出类型。用于混合精度。

返回:

Output – 4 维张量,形状为 [batch, out_height, out_width, out_channel]

返回类型:

tvm.te.Tensor

tvm.topi.nn.group_conv2d_transpose_nchw(data, kernel, stride, padding, out_dtype, output_padding, groups)

NCHW 布局中的分组卷积运算符。

参数:
  • data (tvm.te.Tensor) – 4 维,形状为 [batch, in_channel, in_height, in_width]

  • kernel (tvm.te.Tensor) – 4 维,形状为 [in_channel, out_channel // groups, filter_height, filter_width]

  • stride (int两个整数的列表/元组) – 步幅大小,或 [stride_height, stride_width]

  • padding (int2 个或 4 个整数的列表/元组) – 填充大小,或 2 个整数的 [pad_height, pad_width],或 4 个整数的 [pad_top, pad_left, pad_bottom, pad_right]

  • out_dtype (str) – 输出数据类型。用于混合精度。

  • output_padding (tuple of ints) – 用于获取正确的梯度输出形状

  • groups (int) – 分组数

  • out_dtype – 输出类型。用于混合精度。

返回:

Output – 4-D 张量,形状为 [batch, out_channel, out_height, out_width]

返回类型:

tvm.te.Tensor

tvm.topi.nn.group_conv3d_transpose_ncdhw(data, kernel, strides, padding, out_dtype, output_padding, groups)

转置分组 3D 卷积 ncdhw 前向运算符。

参数:
  • data (tvm.te.Tensor) – 5 维,形状为 [batch, in_channel, in_depth, in_height, in_width]

  • kernel (tvm.te.Tensor) – 5 维,形状为 [in_channel, num_filter, filter_depth, filter_height, filter_width]

  • strides (int一个包含三个整数的列表/元组) – 沿深度、高度和宽度的空间步长

  • padding (intstr) – 填充大小,或 [‘VALID’, ‘SAME’]

  • out_dtype (str) – 输出数据类型。用于混合精度。

  • output_padding (tuple of ints) – 用于获取正确的梯度输出形状

  • groups (int) – 分组数

返回:

Output – 5 维张量,形状为 [batch, out_channel, out_depth, out_height, out_width]

返回类型:

tvm.te.Tensor

tvm.topi.nn.leaky_relu(x, alpha)

对输入 x 执行 leaky relu。

参数:
返回:

y – 结果。

返回类型:

tvm.te.Tensor

tvm.topi.nn.log_softmax(x, axis=-1)

对数据执行 log softmax 激活

参数:

data (tvm.te.Tensor) – N 维输入数据

返回:

output – N 维输出,形状相同

返回类型:

tvm.te.Tensor

tvm.topi.nn.lrn(data, size, axis=1, alpha=0.0001, beta=0.75, bias=2)

对输入数据执行跨通道局部响应归一化。

sum_sqr_up^i{x, y} = (bias+((alpha/size)* {sum_{j=max(0, i-size/2)}^{min(N-1,i+size/2)} (data^j{x,y})^2}))^beta output^i{x, y} = data^i{x, y}/sum_sqr_up^i{x, y} N 是输入通道数

参数:
  • data (tvm.te.Tensor) – 4 维,形状为 [batch, channel, height, width]

  • size (int) – 归一化窗口大小

  • axis (int) – 输入数据布局通道轴,默认值为 1 (NCHW 格式)

  • bias (float) – 避免除以 0 的偏移量

  • alpha (float) – 被除数

  • beta (float) – 指数

返回:

output – 4 维输出,形状相同

返回类型:

tvm.te.Tensor

tvm.topi.nn.lstm(Xs, Wi, Wh, Bi=None, Bh=None, h_init=None, c_init=None, proj=None, p_i=None, p_f=None, p_o=None, f_act=<function sigmoid>, g_act=<function tanh>, h_act=<function tanh>, reverse=False, weight_layout: str = 'IFGO')

使用 TE scan 实现的通用 LSTM。

参数:
  • Xs (te.Tensor) – 输入序列,形状为 (seq_len, batch_size, in_dim)

  • Wi (te.Tensor) – 输入权重矩阵,形状为 (4 * hidden_dim, in_dim)。权重根据 weight_layout 打包。

  • Wh (te.Tensor) – 隐藏层权重矩阵,形状为 (4 * hidden_dim, hidden_dim or proj_dim)。与 Wh 一起打包。

  • Bi (te.Tensor, optional) – 输入偏置,形状为 (4 * hidden_dim,),默认为 None。与 Wh 一起打包。

  • Bh (te.Tensor, optional) – 隐藏层偏置,形状与 Bi 相同,默认为 None。与 Wh 一起打包。

  • h_init (te.Tensor, optional) – 初始隐藏状态,形状为 (batch_size, hidden_dim or proj_dim),如果为 None 则为零

  • c_init (te.Tensor, optional) – 初始细胞状态,形状与 h_init 相同,如果为 None 则为零

  • proj (te.Tensor, optional) – 投影矩阵,形状为 (proj_dim, hidden_dim),默认为 None

  • p_i (te.Tensor, optional) – 窥视孔 LSTM 矩阵,形状为 (batch_size, hidden_dim),默认为 None

  • p_f (te.Tensor, optional) – 窥视孔 LSTM 矩阵,形状为 (batch_size, hidden_dim),默认为 None

  • p_o (te.Tensor, optional) – 窥视孔 LSTM 矩阵,形状为 (batch_size, hidden_dim),默认为 None

  • f_act (F, optional) – 门激活函数

  • g_act (F, optional) – 门激活函数

  • h_act (F, optional) – 门激活函数

  • reverse (bool, optional) – 是否反向处理 Xs,默认为 False

  • weight_layout (str, optional) – 门控的打包权重布局,默认为 “IFGO”。注意:I = 输入, F = 遗忘, G = 细胞, O = 输出。

返回:

result – 隐藏状态(形状为 (seq_len, batch_size, hidden_dim or proj_dim))和细胞状态(形状为 (seq_len, batch_size, hidden_dim))的元组。

返回类型:

te.Tensor, te.Tensor

tvm.topi.nn.matmul(tensor_a, tensor_b, bias=None, out_dtype=None, transpose_a=False, transpose_b=False, auto_scheduler_rewritten_layout='', meta_schedule_original_shape=None)

topi 中 matmul 的默认实现。

参数:
  • tensor_a (tvm.te.Tensor) – 2 维,形状为 [batch, in_dim]

  • tensor_b (tvm.te.Tensor) – 2 维,形状为 [out_dim, in_dim]

  • bias (Optional[tvm.te.Tensor]) – 1 维张量,形状为 [out_dim]

  • out_dtype (Optional[str]) – 输出类型。用于混合精度。

  • transpose_a (Optional[bool] = False) – tensor_a 是否为转置格式。

  • transpose_b (Optional[bool] = False) – tensor_b 是否为转置格式。

  • auto_scheduler_rewritten_layout (Optional[str] = "") – 自动调度器的布局重写传递后的布局。

  • meta_schedule_original_shape (Optional[List[PrimExpr]] = None) – 输入张量的原始形状。

返回:

output – 2 维,形状为 [batch, out_dim]

返回类型:

tvm.te.Tensor

tvm.topi.nn.mirror_pad(data, pad_before, pad_after=None, mode='SYMMETRIC', name='MirrorPadInput')

使用对称或反射镜像填充输入。

参数:
  • data (tvm.te.Tensor) – n 维输入,可以是任何布局。

  • pad_before (list / tuple of n ints) – 在每个维度上填充的宽度,用于在轴开始之前填充。

  • pad_after (list / tuple of n ints, optional) – 在每个维度上填充的宽度,用于在轴结束后填充。

  • mode (str, optional) – 要应用的镜像填充类型。必须为 SYMMETRIC 或 REFLECT

  • name (str, optional) – 生成的运算符的名称前缀

返回:

Output – n 维,布局与 Input 相同。

返回类型:

tvm.te.Tensor

tvm.topi.nn.namedtuple(typename, field_names, *, rename=False, defaults=None, module=None)

返回具有命名字段的元组的新子类。

>>> Point = namedtuple('Point', ['x', 'y'])
>>> Point.__doc__                   # docstring for the new class
'Point(x, y)'
>>> p = Point(11, y=22)             # instantiate with positional args or keywords
>>> p[0] + p[1]                     # indexable like a plain tuple
33
>>> x, y = p                        # unpack like a regular tuple
>>> x, y
(11, 22)
>>> p.x + p.y                       # fields also accessible by name
33
>>> d = p._asdict()                 # convert to a dictionary
>>> d['x']
11
>>> Point(**d)                      # convert from a dictionary
Point(x=11, y=22)
>>> p._replace(x=100)               # _replace() is like str.replace() but targets named fields
Point(x=100, y=22)
tvm.topi.nn.nll_loss(predictions, targets, weights, reduction, ignore_index)

输入数据的负对数似然损失。

output{n, i_1, i_2, …, i_k} = -p * w
where t = target{n, i_1, i_2, …, i_k}

p = predictions{n, t, i_1, i_2, i_k} w = weights{n, i_1, i_2, …, i_k} if t != ignore_index else 0

result = reduction(output)

参数:
  • predictions (tvm.te.Tensor) – (k+2) 维,形状为 (N, C, d_1, d_2, …, d_k),其中 C 是目标类别的数量

  • targets (tvm.te.Tensor) – (k+1) 维,形状为 (N, d_1, d_2, …, d_k) 输入的目标值。

  • weights (tvm.te.Tensor) – 1 维,形状为 (C,) 每个目标值的权重。

  • reduction (string) – 应用于输出的缩减方法。可以是 “mean”(平均值)、“sum”(总和)或 “none”(无)。

  • ignore_index (int) – 要忽略的目标值。

返回:

output – 如果缩减类型为 “mean” 或 “sum”,则为标量;否则与 target 形状相同。

返回类型:

tvm.te.Tensor

tvm.topi.nn.pad(data, pad_before, pad_after=None, pad_value=0.0, name='PadInput', attrs=None)

用零填充输入。

参数:
  • data (tvm.te.Tensor) – n 维输入,可以是任何布局。

  • pad_before (list / tuple of n ints) – 在每个维度上填充的宽度,用于在轴开始之前填充。

  • pad_after (list / tuple of n ints, optional) – 在每个维度上填充的宽度,用于在轴结束后填充。

  • pad_value (float, optional) – 要填充的值。

  • name (str, optional) – 生成的运算符的名称前缀

返回:

Output – n 维,布局与 Input 相同。

返回类型:

tvm.te.Tensor

tvm.topi.nn.pool1d(data, kernel, stride, dilation, padding, pool_type, ceil_mode=False, layout='NCW', count_include_pad=True)
对数据的宽度维度执行池化。

宽度轴根据布局字符串确定。其中 ‘w’ 表示宽度。宽度维度不能被分割。例如,NCW、NCW16c 等对于池化有效,而 NCW16w 无效。有关布局字符串约定的更多信息,请参阅参数 layout

参数:
  • data (tvm.te.Tensor) – n 维,形状为布局

  • kernel (list/tuple of one int or int) – 卷积核大小,[kernel_width]

  • stride (list/tuple of one int or int) – 步幅大小,[stride_width]

  • dilation (list/tuple of two ints) – 空洞大小,[dilation_height, dilation_width]

  • padding (list/tuple of two ints) – 填充大小,[pad_left, pad_right]

  • pool_type (str) – 池化类型,‘max’ 或 ‘avg’

  • ceil_mode (bool) – 在计算输出大小时是否使用 ceil 模式。

  • layout (string) – 输入数据的布局。布局应由大写字母、小写字母和数字组成,其中大写字母表示维度,相应的小写字母和因子大小表示分割维度。例如,NCW16c 可以描述一个 4 维张量 [batch_size, channel, width, channel_block],其中 channel_block=16 是维度 channel 的分割。

  • count_include_pad (bool) – 当 pool_type 为 ‘avg’ 时,是否在计算中包含填充。

返回:

output – n 维,布局相同

返回类型:

tvm.te.Tensor

tvm.topi.nn.pool2d(data, kernel, stride, dilation, padding, pool_type, ceil_mode=False, layout='NCHW', count_include_pad=True)
对数据的高度和宽度维度执行池化。

它根据布局字符串决定高度和宽度维度,其中 ‘W’ 和 ‘H’ 分别表示宽度和高度。宽度和高度维度不能拆分。例如,NCHW、NCHW16c 等对于池化有效,而 NCHW16w、NCHW16h 则无效。有关布局字符串约定的更多信息,请参见 layout 参数。

参数:
  • data (tvm.te.Tensor) – n 维,形状为布局

  • kernel (list/tuple of two ints) – 卷积核大小,[kernel_height, kernel_width]

  • stride (list/tuple of two ints) – 步幅大小,[stride_height, stride_width]

  • dilation (list/tuple of two ints) – 空洞大小,[dilation_height, dilation_width]

  • padding (list/tuple of four ints) – 填充大小,[pad_top, pad_left, pad_bottom, pad_right]]

  • pool_type (str) – 池化类型,‘max’ 或 ‘avg’

  • ceil_mode (bool) – 在计算输出大小时是否使用 ceil 模式。

  • layout (string) – 输入数据的布局。布局应由大写字母、小写字母和数字组成,其中大写字母表示维度,相应的小写字母和因子大小表示拆分维度。例如,NCHW16c 可以描述一个 5 维张量,形状为 [batch_size, channel, height, width, channel_block],其中 channel_block=16 是维度通道的拆分。

  • count_include_pad (bool) – 当 pool_type 为 ‘avg’ 时,是否在计算中包含填充。

返回:

output – n 维,布局相同

返回类型:

tvm.te.Tensor

tvm.topi.nn.pool3d(data, kernel, stride, dilation, padding, pool_type, ceil_mode=False, layout='NCDHW', count_include_pad=True)
对数据的深度、高度和宽度维度执行池化操作。

它根据布局字符串确定深度、高度和宽度维度,其中 ‘D’、‘W’ 和 ‘H’ 分别表示深度、宽度和高度。深度、宽度和高度维度不能被分割。例如,NCDHW、NCDHW16c 等对于池化有效,而 NCDHW16d、NCDHW16w、NCDHW16h 无效。有关布局字符串约定的更多信息,请参阅参数 layout

参数:
  • data (tvm.te.Tensor) – n 维,形状为布局

  • kernel (list/tuple of three ints) – 卷积核大小,[kernel_depth, kernel_height, kernel_width]

  • stride (list/tuple of three ints) – 步幅大小,[stride_depth, stride_height, stride_width]

  • dilation (list/tuple of two ints) – 空洞大小,[dilation_height, dilation_width]

  • padding (list/tuple of six ints) – 填充大小,[pad_front, pad_top, pad_left, pad_back, pad_bottom, pad_right]

  • pool_type (str) – 池化类型,‘max’ 或 ‘avg’

  • ceil_mode (bool) – 在计算输出大小时是否使用 ceil 模式。

  • layout (string) – 输入数据的布局。布局应由大写字母、小写字母和数字组成,其中大写字母表示维度,相应的小写字母和因子大小表示分割维度。例如,NCDHW16c 可以描述一个 6 维张量 [batch_size, channel, depth, height, width, channel_block],其中 channel_block=16 是维度 channel 的分割。

  • count_include_pad (bool) – 当 pool_type 为 ‘avg’ 时,是否在计算中包含填充。

返回:

output – n 维,布局相同

返回类型:

tvm.te.Tensor

tvm.topi.nn.pool_grad(grads, data, kernel, stride, padding, pool_type, ceil_mode=False, count_include_pad=True, layout='NCHW')
数据的高度和宽度维度上的池化梯度。

它根据布局字符串决定高度和宽度维度,其中 ‘W’ 和 ‘H’ 分别表示宽度和高度。宽度和高度维度不能拆分。例如,NCHW、NCHW16c 等对于池化有效,而 NCHW16w、NCHW16h 则无效。有关布局字符串约定的更多信息,请参见 layout 参数。

参数:
  • grads (tvm.te.Tensor) – n 维,形状为布局

  • data (tvm.te.Tensor) – n 维,形状为布局

  • kernel (list/tuple of two ints) – 卷积核大小,[kernel_height, kernel_width]

  • stride (list/tuple of two ints) – 步幅大小,[stride_height, stride_width]

  • padding (list/tuple of four ints) – 填充大小,[pad_top, pad_left, pad_bottom, pad_right]]

  • pool_type (str) – 池化类型,‘max’ 或 ‘avg’

  • ceil_mode (bool) – 在计算输出大小时是否使用 ceil 模式。

  • count_include_pad (bool) – 当 pool_type 为 ‘avg’ 时,是否在计算中包含填充。

  • layout (string) – 输入数据的布局。布局应由大写字母、小写字母和数字组成,其中大写字母表示维度,相应的小写字母和因子大小表示拆分维度。例如,NCHW16c 可以描述一个 5 维张量,形状为 [batch_size, channel, height, width, channel_block],其中 channel_block=16 是维度通道的拆分。

返回:

output – n 维,布局相同

返回类型:

tvm.te.Tensor

tvm.topi.nn.prelu(x, slope, axis=1)

PReLU。它接受两个参数:输入 x 和权重数组 W,并将输出计算为 \(PReLU(x) y = x > 0 ? x : W * x\),其中 \(*\) 是批次中每个样本的元素级乘法。

参数:
返回:

  • y (tvm.te.Tensor) – 结果。

  • 链接

  • —–

  • [http (//arxiv.org/pdf/1502.01852v1.pdf])

tvm.topi.nn.reduce(function, sequence[, initial]) value

从左到右将一个接受两个参数的函数累积地应用于序列的项,从而将序列归约为单个值。例如,reduce(lambda x, y: x+y, [1, 2, 3, 4, 5]) 计算 ((((1+2)+3)+4)+5)。如果提供了 initial,它将在计算中放在序列项之前,并在序列为空时作为默认值。

tvm.topi.nn.relu(x)

对输入 x 执行 ReLU 运算。

参数:

x (tvm.te.Tensor) – 输入参数。

返回:

y – 结果。

返回类型:

tvm.te.Tensor

tvm.topi.nn.scale_shift_nchw(Input, Scale, Shift)

推理中的批归一化运算符。

参数:
  • Input (tvm.te.Tensor) – 4-D 输入张量,NCHW 布局 [batch, channel, height, width]

  • Scale (tvm.te.Tensor) – 缩放张量,1-D,大小为通道数

  • Shift (tvm.te.Tensor) – 偏移张量,1-D,大小为通道数

返回:

Output – 输出张量,布局为 NCHW

返回类型:

tvm.te.Tensor

tvm.topi.nn.scale_shift_nchwc(Input, Scale, Shift)

推理中的批归一化运算符。

参数:
  • Input (tvm.te.Tensor) – 5-D 输入张量,NCHWc 布局 [batch, channel_chunk, height, width, channel_block]

  • Scale (tvm.te.Tensor) – 缩放张量,2-D,大小为 [channel_chunk, channel_block]

  • Shift (tvm.te.Tensor) – 偏移张量,2-D,大小为 [channel_chunk, channel_block]

返回:

Output – 输出张量,布局为 NHWC

返回类型:

tvm.te.Tensor

tvm.topi.nn.scale_shift_nhwc(Input, Scale, Shift)

推理中的批归一化运算符。

参数:
  • Input (tvm.te.Tensor) – 4-D 输入张量,NHWC 布局 [batch, height, width, channel]

  • Scale (tvm.te.Tensor) – 缩放张量,1-D,大小为通道数

  • Shift (tvm.te.Tensor) – 偏移张量,1-D,大小为通道数

返回:

Output – 输出张量,布局为 NHWC

返回类型:

tvm.te.Tensor

tvm.topi.nn.simplify(expr)

如果表达式是 Expr,则简化表达式;如果表达式是 int,则直接返回。

参数:

expr (Expr or int) – 输入。

返回:

out – 简化的输出

返回类型:

Expr or int

tvm.topi.nn.simulated_dequantize(data, in_dtype, input_scale=None, input_zero_point=None, axis=-1)

模拟 QNN 反量化运算符,模拟 QNN 输出,但不更改数据类型。与真正的 QNN 反量化相比,此运算符的优势在于它允许动态数据类型选择,并且可以对每个通道和标量比例以及零点进行操作,而 QNN 反量化要求这些在编译时都是固定的。

参数:
  • data (tvm.te.Tensor) – 运算符的 N-D 输入张量。

  • in_dtype (tvm.te.Tensor) – 一个标量变量,指示要模拟反量化使用的数据类型。使用 SQNN_DTYPE_TO_CODE 将 dtype 字符串转换为相应的变量值。

  • input_scale (tvm.te.Tensor, optional) – 一个标量张量,表示从整数数据类型反量化时使用的比例。当它包含多个值时,N 必须与数据中的通道数匹配。

  • input_zero_point (tvm.te.Tensor, optional) – 一个 1-D 张量,表示从整数数据类型反量化时使用的零点。当它包含多个值时,N 必须与数据中的通道数匹配。

  • axis (int, optional) – 量化的通道轴。默认值为 -1,对应于最后一个轴。

tvm.topi.nn.simulated_quantize(data, out_dtype, output_scale=None, output_zero_point=None, axis=-1)

模拟 QNN 量化运算符,模拟 QNN 输出,但不更改数据类型。与真正的 QNN 量化相比,此运算符的优势在于它允许动态数据类型选择,并且可以对每个通道和标量比例以及零点进行操作,而 QNN 量化要求这些在编译时都是固定的。

参数:
  • data (tvm.te.Tensor) – 运算符的 N-D 输入张量。

  • out_dtype (tvm.te.Tensor) – 一个标量变量,指示要模拟量化使用的数据类型。使用 SQNN_DTYPE_TO_CODE 将 dtype 字符串转换为相应的变量值。

  • output_scale (tvm.te.Tensor, optional) – 一个标量张量,表示量化为整数数据类型时使用的比例。当它包含多个值时,N 必须与数据中的通道数匹配。

  • output_zero_point (tvm.te.Tensor, optional) – 一个 1-D 张量,表示量化为整数数据类型时使用的零点。当它包含多个值时,N 必须与数据中的通道数匹配。

  • axis (int, optional) – 量化的通道轴。默认值为 -1,对应于最后一个轴。

tvm.topi.nn.softmax(x, axis=-1)

对数据执行 softmax 激活。

参数:
返回:

output – 输出形状与输入相同

返回类型:

tvm.te.Tensor

tvm.topi.nn.softmax_common(x, axis, use_fast_exp)

softmax 和 fast_softmax 的公共部分

tvm.topi.nn.space_to_batch_nd(data, block_shape, pad_before, pad_after, pad_value=0.0)

对数据执行批次到空间转换

参数:
  • data (tvm.te.Tensor) – N 维张量,形状为 [batch, spatial_shape, remaining_shapes],其中 spatial_shape 具有 M 个维度。

  • block_shape (list of ints) – 大小为 [M] 的列表,其中 M 是空间维度数,指定每个空间维度的块大小。

  • pad_before (list of ints) – 形状为 [M] 的列表,其中 M 是空间维度数,指定每个空间维度之前的零填充大小。

  • pad_after (list of ints) – 形状为 [M] 的列表,其中 M 是空间维度数,指定每个空间维度之后的零填充大小。

  • pad_value (float, optional) – 用于填充的值。

返回:

output

返回类型:

tvm.te.Tensor

tvm.topi.nn.space_to_depth(data, block_size, layout='NCHW')

对数据执行空间到深度转换

参数:
  • data (tvm.te.Tensor) – 4 维张量,布局为 NCHW 或 NHWC。

  • block_size (int) – 分解为通道维度的块大小。

  • layout (string) – NCHW 或 NHWC,表示数据布局。

返回:

output – 形状为 [N, C * block_size**2, H / block_size, W / block_size] 的输出

返回类型:

tvm.te.Tensor

tvm.topi.nn.strided_slice(a, begin, end, strides=None, axes=None, slice_mode='end', assume_inbound=True)

数组切片。

参数:
  • a (tvm.te.Tensor) – 要切片的张量。

  • begin (list of int) – 切片开始的索引。

  • end (list of int) – 指示切片结束的索引。

  • strides (list of int, optional) – 指定步长值,可以是负数,在这种情况下,输入张量将在该特定轴上反转。

  • axes (list of int, optional) – 应用切片的轴。当指定时,begin、end、strides 和 axes 需要是相同长度的整数列表。

  • slice_mode (str, optional) – 切片模式 [end, size]。end - 切片的结束索引 [默认]。size - 输入步长将被忽略,此模式下的输入 end 指示从 begin 指定的位置开始的切片的大小。如果 end[i] 为 -1,则该维度中的所有剩余元素都包含在切片中。

  • assume_inbound (bool, optional) – 一个标志,指示是否假定所有索引都在界内

返回:

ret

返回类型:

tvm.te.Tensor

tvm.topi.nn.unpack_NCHWc_to_nchw(packed_out, out_dtype)

将 conv2d_NCHWc 输出从 NCHWc 布局解包到 NCHW

参数:
  • packed_out (tvm.te.Tensor) – conv2d_NCHWc 的输出张量。

  • out_dtype (str) – 输出数据类型。

返回:

unpacked_out – NCHW 布局中解包的输出张量。

返回类型:

tvm.te.Tensor

tvm.topi.nn.upsampling(data, scale_h, scale_w, layout='NCHW', method='nearest_neighbor', align_corners=False, output_shape=None)
对数据执行上采样。

支持最近邻和双线性上采样。

参数:
  • inputs (tvm.te.Tensor) – inputs 是一个 4-D 张量,形状为 [batch, channel, in_height, in_width] 或 [batch, in_height, in_width, channel]

  • scale_h (float) – 高度的缩放因子

  • scale_w (float) – 宽度的缩放因子

  • layout (string, optional) – “NCHW” 或 “NHWC”

  • method ({"bilinear", "nearest_neighbor", "bicubic"}) – 用于上采样的方法。

  • output_shape (tvm.tir.container.Array, optional) – 要返回的形状。如果留空,将推断形状(如果形状是动态确定的,则传递 out_dtype.shape 作为 output_shape)

返回:

output – 4-D,形状为 [batch, channel, in_height*scale_h, in_width*scale_w] 或 [batch, in_height*scale, in_width*scale, channel]

返回类型:

tvm.te.Tensor

tvm.topi.nn.upsampling3d(data, scale_d, scale_h, scale_w, layout='NCDHW', method='nearest_neighbor', coordinate_transformation_mode='half_pixel', output_shape=None)
对数据执行上采样。

支持最近邻和双线性上采样。

参数:
  • inputs (tvm.te.Tensor) – inputs 是一个 5-D 张量,形状为 [batch, channel, in_depth, in_height, in_width] 或 [batch, in_depth, in_height, in_width, channel]

  • scale_d (float) – 深度的缩放因子

  • scale_h (float) – 高度的缩放因子

  • scale_w (float) – 宽度的缩放因子

  • layout (string, optional) – “NCDHW” 或 “NDHWC”

  • method ({"trilinear", "nearest_neighbor"}) – 用于上采样的方法。

  • coordinate_transformation_mode (string, optional) – 描述如何将调整大小的张量中的坐标转换为原始张量中的坐标。有关详细信息,请参阅 ONNX Resize 运算符规范。可用选项为 “half_pixel”、“align_corners” 和 “asymmetric”。

  • output_shape (tvm.tir.container.Array, optional) – 要返回的形状。如果留空,将推断形状(如果形状是动态确定的,则传递 out_dtype.shape 作为 output_shape)

返回:

output – 5-D,形状为 [batch, channel, in_depth*scale, in_height*scale, in_width*scale] 或 [batch, in_depth*scale, in_height*scale, in_width*scale, channel]

返回类型:

tvm.te.Tensor

tvm.topi.nn.winograd_transform_matrices(tile_size, kernel_size, out_dtype)

计算 tile_size 的 A、B 和 G 变换矩阵,结果为 tvm.Expr

tvm.topi.nn.instance_norm(data, gamma, beta, axis, epsilon=1e-05)

实例归一化运算符。

参数:
  • data (tvm.te.Tensor) – N-D,形状为 (d_0, d_1, …, d_{N-1})

  • gamma (tvm.te.Tensor) – K-D,形状为 (r_0, r_1, …, r_{K-1}),其中 K == len(axis) 且 d_{axis_k} == r_k

  • beta (tvm.te.Tensor) – 可选,K-D,形状为 (r_0, r_1, …, r_{K-1}),其中 K == len(axis) 且 d_{axis_k} == r_k

  • axis (list of int) – 应用归一化的轴(计算均值和方差的轴)

  • epsilon (float) – epsilon 值,用于避免被零除。

返回:

result – N-D,形状为 (d_0, d_1, …, d_{N-1})

返回类型:

tvm.te.Tensor

tvm.topi.nn.layer_norm(data, gamma, beta, axis, epsilon=1e-05)

层归一化运算符。它接受 fp16 和 fp32 作为输入数据类型。它会将输入转换为 fp32 以执行计算。输出将具有与输入相同的数据类型。

参数:
  • data (tvm.te.Tensor) – N-D,形状为 (d_0, d_1, …, d_{N-1})

  • gamma (tvm.te.Tensor) – K-D,形状为 (r_0, r_1, …, r_{K-1}),其中 K == len(axis) 且 d_{axis_k} == r_k

  • beta (tvm.te.Tensor) – 可选,K-D,形状为 (r_0, r_1, …, r_{K-1}),其中 K == len(axis) 且 d_{axis_k} == r_k

  • axis (list of int) – 应用归一化的轴

  • epsilon (float) – epsilon 值,用于避免被零除。

返回:

result – N-D,形状为 (d_0, d_1, …, d_{N-1})

返回类型:

tvm.te.Tensor

tvm.topi.nn.group_norm(data, gamma, beta, num_groups, channel_axis, axes, epsilon=1e-05)

组归一化运算符。它接受 fp16 和 fp32 作为输入数据类型。它会将输入转换为 fp32 以执行计算。输出将具有与输入相同的数据类型。

参数:
  • data (tvm.te.Tensor) – N-D,形状为 (d_0, d_1, …, d_{N-1})

  • gamma (tvm.te.Tensor) – 1-D,形状为 (r_0),其中 r_0 == d_{channel_axis}

  • beta (tvm.te.Tensor) – 可选,1-D,形状为 (r_0),其中 r_0 == d_{channel_axis}

  • num_groups (int) – 组的数量

  • channel_axis (int) – 通道轴

  • axes (list of int) – 应用归一化的轴,不包括通道轴

  • epsilon (float) – epsilon 值,用于避免被零除。

返回:

result – N-D,形状为 (d_0, d_1, …, d_{N-1})

返回类型:

tvm.te.Tensor

tvm.topi.nn.rms_norm(data, weight, axis, epsilon=1e-05)

均方根归一化运算符。输出将具有与输入相同的数据类型。

参数:
  • data (tvm.te.Tensor) – N-D,形状为 (d_0, d_1, …, d_{N-1})

  • weight (tvm.te.Tensor) – K-D,形状为 (r_0, r_1, …, r_{K-1}),其中 K == len(axis) 且 d_{axis_k} == r_k

  • axis (list of int) – 应用归一化的轴

  • epsilon (float) – epsilon 值,用于避免被零除。

返回:

result – N-D,形状为 (d_0, d_1, …, d_{N-1})

返回类型:

tvm.te.Tensor

tvm.topi.image

图像网络运算符

函数

affine_grid(data, target_shape)

affine_grid 运算符,用于生成 2D 采样网格。

can_convert_multiply_to_intdiv(origin_size, ...)

检查是否可以将乘法转换为除法

crop_and_resize(data, boxes, box_indices, ...)

对数据执行裁剪和调整大小操作。

dilation2d_nchw(input, filter, stride, ...)

NCHW 布局中的形态学膨胀运算符。

dilation2d_nhwc(input, filter, stride, ...)

形态学 2d 膨胀 NHWC 布局。

get_1d_indices(indices[, layout])

获取 1d 索引

get_1d_pixel(data, layout, image_width, n, ...)

获取 1d 像素

get_2d_indices(indices[, layout])

获取 2d 索引

get_2d_pixel(data, layout, image_height, ...)

获取 2d 像素

get_3d_indices(indices[, layout])

获取 3d 索引

get_3d_pixel(data, layout, image_depth, ...)

获取 3d 像素

get_closest_index(in_x, rounding_method, boxes)

根据某种舍入方法获取最接近某个值的索引

get_inx(x, image_width, target_width, ...[, ...])

使用各种坐标变换方法从输出 x 推断输入 x

get_pad_tuple(padding, kernel)

获取填充选项的通用代码

grid_sample(data, grid[, method, layout, ...])

将网格采样应用于输入特征图。

nchw_pack_layout(layout_info)

检查布局类型是否为 NCHWinic

nchw_xc_layout(layout_info)

检查布局类型是否为 NCHWxc

pad(data, pad_before[, pad_after, ...])

用零填充输入。

resize1d(data, roi, size[, layout, method, ...])

对数据执行调整大小操作。

resize2d(data, roi, size[, layout, method, ...])

对数据执行调整大小操作。

resize3d(data, roi, size[, layout, method, ...])

对数据执行调整大小操作。

simplify(expr)

如果表达式是 Expr,则简化表达式;如果表达式是 int,则直接返回。

tvm.topi.image.affine_grid(data, target_shape)

affine_grid 运算符,用于生成 2D 采样网格。

此操作在 https://arxiv.org/pdf/1506.02025.pdf 中描述。它在目标形状内生成均匀采样网格,并将其归一化为 [-1, 1]。然后将提供的仿射变换应用于采样网格。

参数:
  • data (tvm.Tensor) – 3-D,形状为 [batch, 2, 3]。仿射矩阵。

  • target_shape (list/tuple of two int) – 指定输出形状 (H, W)。

返回:

Output – 4-D,形状为 [batch, 2, target_height, target_width]

返回类型:

tvm.Tensor

tvm.topi.image.can_convert_multiply_to_intdiv(origin_size, scaled_size)

检查是否可以将乘法转换为除法

tvm.topi.image.crop_and_resize(data, boxes, box_indices, crop_size, layout='NCHW', method='bilinear', extrapolation_value=None, out_dtype=None)

对数据执行裁剪和调整大小操作。

参数:
  • data (tvm.te.Tensor) – inputs 是一个 4-D 张量,形状为 [batch, channel, in_height, in_width] 或 [batch, in_height, in_width, channel]

  • boxes (tvm.te.Tensor) – 形状为 [num_boxes, 4] 的 2-D 张量。张量的每一行指定一个框的坐标。

  • box_indices (tvm.te.Tensor) – 形状为 [num_boxes] 的 1-D 张量,box_indices[i] 指定第 i 个框引用的数据。

  • crop_size (Tuple) – 每个框的目标大小。

  • layout (string, optional) – “NCHW”、“NHWC”

  • method ({"bilinear", "nearest_neighbor"}) – 用于调整大小的方法。

  • extrapolation_value (float, 可选) – 适用时,用于外推的值。

  • out_dtype (string, 可选) – 返回类型。如果留空,则与输入类型相同。

返回:

output – 4-D 张量,形状为 [num_boxes, channel, crop_height, crop_width] 或 [num_boxes, crop_height, crop_width, channel]

返回类型:

tvm.te.Tensor

tvm.topi.image.dilation2d_nchw(input, filter, stride, padding, dilations, out_dtype=None)

NCHW 布局中的形态学膨胀运算符。

参数:
  • input (tvm.te.Tensor) – 4-D 张量,形状为 [batch, in_channel, in_height, in_width]

  • filter (tvm.te.Tensor) – 3-D 张量,形状为 [ in_channel, filter_height, filter_width]

  • stride (int两个整数的列表/元组) – 步幅大小,或 [stride_height, stride_width]

  • padding (intstr) – 填充大小

  • dilations (int一个 list/tuple 包含 两个 ints) – 空洞大小,或 [dilation_height, dilation_width]

  • out_dtype (Optional[str]) – 指定输出数据类型。

返回:

Output – 4-D 张量,形状为 [batch, in_channel, out_height, out_width]

返回类型:

tvm.te.Tensor

tvm.topi.image.dilation2d_nhwc(input, filter, stride, padding, dilations, out_dtype=None)

形态学 2d 膨胀 NHWC 布局。

参数:
  • input (tvm.te.Tensor) – 4-D 张量,形状为 [batch, in_height, in_width, in_channel]

  • filter (tvm.te.Tensor) – 3-D 张量,形状为 [filter_height, filter_width, in_channel]

  • stride (int两个整数的列表/元组) – 步幅大小,或 [stride_height, stride_width]

  • padding (int) – 填充大小

  • dilations (int一个 list/tuple 包含 两个 ints) – 空洞大小,或 [dilation_height, dilation_width]

  • out_dtype (Optional[str]) – 指定输出数据类型。

返回:

Output – 4-D 张量,形状为 [batch, out_height, out_width, in_channel]

返回类型:

tvm.te.Tensor

tvm.topi.image.get_1d_indices(indices, layout='NCW')

获取 1d 索引

tvm.topi.image.get_1d_pixel(data, layout, image_width, n, c, x, cc, ib, ic)

获取 1d 像素

tvm.topi.image.get_2d_indices(indices, layout='NCHW')

获取 2d 索引

tvm.topi.image.get_2d_pixel(data, layout, image_height, image_width, n, c, y, x, cc, ib, ic)

获取 2d 像素

tvm.topi.image.get_3d_indices(indices, layout='NCDHW')

获取 3d 索引

tvm.topi.image.get_3d_pixel(data, layout, image_depth, image_height, image_width, n, c, z, y, x, cc)

获取 3d 像素

tvm.topi.image.get_closest_index(in_x, rounding_method, boxes, use_int_div=False)

根据某种舍入方法获取最接近某个值的索引

tvm.topi.image.get_inx(x, image_width, target_width, coordinate_transformation_mode, start_x=0, end_x=-1, use_int_div=False)

使用各种坐标变换方法从输出 x 推断输入 x

tvm.topi.image.get_pad_tuple(padding, kernel)

获取填充选项的通用代码

参数:
  • padding (intstr) – 填充大小,或 [‘VALID’, ‘SAME’]

  • kernel (tuple of int) – 卷积核大小

返回:

  • pad_top (int) – 顶部的填充大小

  • pad_left (int) – 左侧的填充大小

  • pad_down (int) – 底部的填充大小。

  • pad_right (int) – 右侧的填充大小。

tvm.topi.image.grid_sample(data, grid, method='bilinear', layout='NCHW', padding_mode='zeros', align_corners=True)

将网格采样应用于输入特征图。

给定 \(data\)\(grid\),对于 4-D 张量,输出计算方式如下:

\[x_{src} = grid[batch, 0, y_{dst}, x_{dst}] \ y_{src} = grid[batch, 1, y_{dst}, x_{dst}] \ output[batch, channel, y_{dst}, x_{dst}] = G(data[batch, channel, y_{src}, x_{src}])\]

\(x_{dst}\), \(y_{dst}\) 枚举 \(output\) 中的所有空间位置,\(G()\) 表示插值函数。

如果 padding_mode 为 “zeros”,则边界外的点将用零填充;如果 padding_mode 为 “border”,则用边界像素值填充;如果 padding_mode 为 “reflection”,则用内部像素值填充。

如果 align_corners 为 “True”,则 grid 中的左上角 (-1, -1) 和右下角 (1, 1) 将映射到数据的 (0, 0) 和 (h - 1, w - 1);如果 align_corners 为 “False”,则映射到数据的 (-0.5, -0.5) 和 (h - 0.5, w - 0.5)。

输出的形状将为 4-D (data.shape[0], data.shape[1], grid.shape[2], grid.shape[3]) 或 5-D (data.shape[0], data.shape[1], grid.shape[2], grid.shape[3], grid.shape[4])。

该运算符假定 \(grid\) 已被归一化到 [-1, 1]。

grid_sample 通常与 affine_grid 配合使用,affine_grid 生成用于 grid_sample 的采样网格。

参数:
  • data (tvm.Tensor) – 4-D 张量,形状为 [batch, in_channel, in_height, in_width],或 5-D 张量,形状为 [batch, in_channel, in_depth, in_height, in_width]

  • grid (tvm.Tensor) – 4-D 张量,形状为 [batch, 2, out_height, out_width],或 5-D 张量,形状为 [batch, 3, out_depth, out_height, out_width]

  • method (str) – 插值方法,4-D 支持 “nearest”、“bilinear”、“bicubic”,5-D 支持 “nearest”、“bilinear”(“trilinear”)。

  • layout (str) – 输入数据和输出的布局。

  • padding_mode (str) – 用于网格外部值的填充模式,支持 “zeros”、“border”、“reflection”。

  • align_corners (bool) – 从几何角度来看,我们将输入的像素视为正方形而不是点。如果设置为 “True”,则极值(“-1” 和 “1”)被视为指代输入角像素的中心点。如果设置为 “False”,则它们被视为指代输入角像素的角点,从而使采样对分辨率更不敏感。

返回:

Output – 4-D 张量,形状为 [batch, in_channel, out_height, out_width],或 5-D 张量,形状为 [batch, in_channel, out_depth, out_height, out_width]

返回类型:

tvm.Tensor

tvm.topi.image.nchw_pack_layout(layout_info)

检查布局类型是否为 NCHWinic

tvm.topi.image.nchw_xc_layout(layout_info)

检查布局类型是否为 NCHWxc

tvm.topi.image.pad(data, pad_before, pad_after=None, pad_value=0.0, name='PadInput', attrs=None)

用零填充输入。

参数:
  • data (tvm.te.Tensor) – n 维输入,可以是任何布局。

  • pad_before (list / tuple of n ints) – 在每个维度上填充的宽度,用于在轴开始之前填充。

  • pad_after (list / tuple of n ints, optional) – 在每个维度上填充的宽度,用于在轴结束后填充。

  • pad_value (float, optional) – 要填充的值。

  • name (str, optional) – 生成的运算符的名称前缀

返回:

Output – n 维,布局与 Input 相同。

返回类型:

tvm.te.Tensor

tvm.topi.image.resize1d(data, roi, size, layout='NCW', method='linear', coordinate_transformation_mode='half_pixel', rounding_method='', bicubic_alpha=-0.5, bicubic_exclude=0, extrapolation_value=0.0, out_dtype=None, output_shape=None)

对数据执行调整大小操作。

参数:
  • data (tvm.te.Tensor) – 输入是 3-D 张量,形状为 [batch, channel in_width] 或 [batch in_width, channel]

  • roi (Tuple of FloatExpr) – 用于裁剪输入图像的感兴趣区域。 预期大小为 2,格式为 [start_w, end_w]。 仅当 coordinate_transformation_mode 为 tf_crop_and_resize 时使用。

  • size (Tuple) – 输出分辨率缩放至

  • layout (string, 可选) – “NCW”、“NWC” 或 “NCWc”。

  • coordinate_transformation_mode (string, optional) – 描述如何将调整大小的张量中的坐标转换为原始张量中的坐标。有关详细信息,请参阅 ONNX Resize 运算符规范。可用选项为 “half_pixel”、“align_corners” 和 “asymmetric”。

  • method (string, 可选) – 插值方法 (“nearest”、“linear”、“bicubic”)

  • coordinate_transformation_mode – 描述如何将调整大小后的张量中的坐标转换为原始张量中的坐标。[half_pixel、align_corners、asymmetric、pytorch_half_pixel、tf_half_pixel_for_nn 和 tf_crop_and_resize]。

  • rounding_method – 坐标位置的舍入方法

  • bicubic_alpha (float, 可选) – 双三次样条系数

  • bicubic_exclude (bool, 可选:) – 排除图像外部的值以进行双三次插值

  • extrapolation_value (float, 可选) – 适用时,用于外推的值。

  • out_dtype (string, 可选) – 返回类型。如果留空,则与输入类型相同。

  • output_shape (tvm.tir.container.Array, optional) – 要返回的形状。如果留空,将推断形状(如果形状是动态确定的,则传递 out_dtype.shape 作为 output_shape)

返回:

output – 4-D 张量,形状为 [batch, chananel, in_width*scale] 或 [batch, in_width*scale, channel] 或 5-D 张量,形状为 [batch, channel-major, in_width*scale, channel-minor]

返回类型:

tvm.te.Tensor

tvm.topi.image.resize2d(data, roi, size, layout='NCHW', method='linear', coordinate_transformation_mode='half_pixel', rounding_method='', bicubic_alpha=-0.5, bicubic_exclude=0, extrapolation_value=0.0, out_dtype=None, output_shape=None)

对数据执行调整大小操作。

参数:
  • data (tvm.te.Tensor) – inputs 是一个 4-D 张量,形状为 [batch, channel, in_height, in_width] 或 [batch, in_height, in_width, channel]

  • roi (Tuple of FloatExpr) – 用于裁剪输入图像的感兴趣区域。 预期大小为 4,格式为 [start_h, start_w, end_h, end_w]。 仅当 coordinate_transformation_mode 为 tf_crop_and_resize 时使用。

  • size (Tuple) – 输出分辨率缩放至

  • layout (string, 可选) – “NCHW”、“NHWC” 或 “NCHWc”。

  • method (string, 可选) – 插值方法 (“nearest”、“linear”、“bicubic”)

  • coordinate_transformation_mode (string, 可选) – 描述如何将调整大小后的张量中的坐标转换为原始张量中的坐标。[half_pixel、align_corners、asymmetric、pytorch_half_pixel、tf_half_pixel_for_nn 和 tf_crop_and_resize]。

  • rounding_method – 坐标位置的舍入方法

  • bicubic_alpha (float, 可选) – 双三次样条系数

  • bicubic_exclude (bool, 可选:) – 排除图像外部的值以进行双三次插值

  • extrapolation_value (float, 可选) – 适用时,用于外推的值。

  • out_dtype (string, 可选) – 返回类型。如果留空,则与输入类型相同。

  • output_shape (tvm.tir.container.Array, optional) – 要返回的形状。如果留空,将推断形状(如果形状是动态确定的,则传递 out_dtype.shape 作为 output_shape)

返回:

output – 4-D 张量,形状为 [batch, channel, in_height*scale, in_width*scale] 或 [batch, in_height*scale, in_width*scale, channel] 或 5-D 张量,形状为 [batch, channel-major, in_height*scale, in_width*scale, channel-minor]

返回类型:

tvm.te.Tensor

tvm.topi.image.resize3d(data, roi, size, layout='NCDHW', method='linear', coordinate_transformation_mode='half_pixel', rounding_method='', bicubic_alpha=-0.5, bicubic_exclude=0, extrapolation_value=0.0, out_dtype=None, output_shape=None)

对数据执行调整大小操作。

参数:
  • data (tvm.te.Tensor) – 输入是 5-D 张量,形状为 [batch, channel, in_depth, in_height, in_width] 或 [batch, in_depth, in_height, in_width, channel]

  • roi (Tuple of FloatExpr) – 用于裁剪输入图像的感兴趣区域。 预期大小为 6,格式为 [start_d, start_h, start_w, end_d, end_h, end_w]。 仅当 coordinate_transformation_mode 为 tf_crop_and_resize 时使用。

  • size (Tuple) – 输出分辨率缩放至

  • layout (string, 可选) – “NCDHW”、“NDHWC” 或 “NCDHWc”。

  • method (string, 可选) – 插值方法 (“nearest”、“linear”、“bicubic”)

  • coordinate_transformation_mode (string, 可选) – 描述如何将调整大小后的张量中的坐标转换为原始张量中的坐标。[half_pixel、align_corners、asymmetric、pytorch_half_pixel、tf_half_pixel_for_nn 和 tf_crop_and_resize]。

  • rounding_method – 坐标位置的舍入方法

  • bicubic_alpha (float, 可选) – 双三次样条系数

  • bicubic_exclude (bool, 可选:) – 排除图像外部的值以进行双三次插值

  • extrapolation_value (float, 可选) – 适用时,用于外推的值。

  • out_dtype (string, 可选) – 返回类型。如果留空,则与输入类型相同。

  • output_shape (tvm.tir.container.Array, optional) – 要返回的形状。如果留空,将推断形状(如果形状是动态确定的,则传递 out_dtype.shape 作为 output_shape)

返回:

output – 4-D 张量,形状为 [batch, channel, in_depth*scale, in_height*scale, in_width*scale] 或 [batch, in_depth*scale, in_height*scale, in_width*scale, channel] 或 5-D 张量,形状为 [batch, channel-major, in_depth*scale, in_height*scale, in_width*scale, channel-minor]

返回类型:

tvm.te.Tensor

tvm.topi.image.simplify(expr)

如果表达式是 Expr,则简化表达式;如果表达式是 int,则直接返回。

参数:

expr (Expr or int) – 输入。

返回:

out – 简化的输出

返回类型:

Expr or int