CANN/ops-math矩阵对角线生成算子

MatrixDiagV3

【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库，实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math

产品支持情况

功能说明

• 算子功能：根据输入的单条或多条对角线值生成矩阵，对角线带之外的位置使用$\mathrm{padding_value}$填充。
• 设输出张量最后两维大小分别为$\mathrm{num_rows}$和$\mathrm{num_cols}$，$k=[k_l, k_u]$表示待写入的对角线范围，单对角线场景下$k_l = k_u$。令$d = j - i$表示位置$(i, j)$所在的对角线编号，则最大对角线长度为：

$$ \mathrm{max_diag_len} = \min(\mathrm{num_rows} + \min(k_u, 0),\ \mathrm{num_cols} - \max(k_l, 0)) $$

输出元素满足：

$$ y_{..., i, j} = \begin{cases} x_{..., k_u - d,\ p(i, j)} & k_l \le d \le k_u \ \text{padding_value} & \text{otherwise} \end{cases} $$

其中，$p(i, j)$表示对角线元素在输入$x$最后一维中的位置，具体由属性align控制左右对齐方式。当$k$为单个整数或$k_l = k_u$时，$x$表示单条对角线；当$k_l < k_u$时，$x$表示对角线带。

参数说明

约束说明

• x的秩至少为1。
• k的元素个数只能为1或2；当k为2个元素时，必须满足k[0] <= k[1]。
• 当k表示对角线带时，x的倒数第二维长度必须等于$k_u - k_l + 1$，最后一维长度必须等于$\mathrm{max_diag_len}$。

调用说明

【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库，实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math