企业官网建设流程全解析

pypto.axpy_

【免费下载链接】pyptoPyPTO（发音: pai p-t-o）：Parallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pypto

产品支持情况

功能说明

执行 AXPY 操作：y = alpha * x + y。该操作对 y 张量进行原地更新。

计算公式如下：

$$ y_i = \alpha \cdot x_i + y_i $$

重要说明：AXPY 是原地操作，y 张量会被直接修改。如果后续计算需要使用执行 AXPY 之前的原始 y 值，请在调用 AXPY 前使用pypto.clone(y)进行备份。

函数原型

axpy_(y: Tensor, x: Tensor, alpha: Union[int, float] = 1.0) -> Tensor

参数说明

返回值说明

返回更新后的 y 张量（与输入 y 共享同一内存地址），Tensor 的数据类型与 y 相同，Shape 与 y 相同。

约束说明

1. dtype 约束：
   • 相同 dtype：支持 DT_FP32 + DT_FP32、DT_FP16 + DT_FP16、DT_BF16 + DT_BF16。
   • 混合 dtype：仅支持 DT_FP32 (y) + DT_FP16 (x)，其他组合不支持。
   • DT_BF16 在内部会转换为 DT_FP32 进行计算。
2. 广播约束：
   • y 张量不支持广播。如果 y 的某个维度为 1 而 x 对应维度不为 1，将报错。
   • x 张量支持广播到 y 的形状。
3. Shape 约束：y 和 x 的维度数必须相同（1-4维）。
4. Format 约束：y 和 x 的 Format 必须一致。
5. 原地更新注意：由于 AXPY 是原地操作，y 的原始值会被覆盖。如需保留原始 y 值，请提前 clone：

y_backup = pypto.clone(y) # 备份原始 y 值 y.axpy_(x, alpha=2.0) # y 被原地更新 # 此时 y_backup 仍保留原始值，可用于后续计算

调用示例

TileShape设置示例

调用该 operation 接口前，应通过set_vec_tile_shapes设置 TileShape。

TileShape 维度应和输出一致。

示例：输入 y shape 为[m, n]，x shape 为[m, n]（或[m, 1]广播场景），输出 shape 为[m, n]，TileShape 设置为[m1, n1]，则m1,n1分别用于切分输出的m,n轴。

pypto.set_vec_tile_shapes(32, 32)

接口调用示例

基本用法

y = pypto.tensor([1, 3], pypto.DT_FP32) x = pypto.tensor([1, 3], pypto.DT_FP32) y.axpy_(x, alpha=2.0)

结果示例如下：

输入数据 y: [[1.0 2.0 3.0]] 输入数据 x: [[2.0 3.0 4.0]] alpha: 2.0 输出数据 y: [[5.0 8.0 11.0]] # y = 2.0 * x + y

广播场景

y = pypto.tensor([64, 64], pypto.DT_FP32) # y shape: [64, 64] x = pypto.tensor([64, 1], pypto.DT_FP32) # x shape: [64, 1] (广播到 [64, 64]) y.axpy_(x, alpha=1.5)

保留原始 y 值

y = pypto.tensor([32, 32], pypto.DT_FP32) x = pypto.tensor([32, 32], pypto.DT_FP32) # 如需使用原始 y 值，提前备份 y_backup = pypto.clone(y) # 执行 AXPY，y 被原地更新 y.axpy_(x, alpha=2.0) # y_backup 仍保留原始值，可用于其他计算 diff = pypto.sub(y, y_backup) # 计算 y 与原始值的差值

混合精度（FP32 + FP16）

y = pypto.tensor([32, 32], pypto.DT_FP32) # y 为 FP32 x = pypto.tensor([32, 32], pypto.DT_FP16) # x 为 FP16 y.axpy_(x, alpha=1.0) # 支持 FP32(y) + FP16(x)

BF16 dtype

y = pypto.tensor([32, 32], pypto.DT_BF16) x = pypto.tensor([32, 32], pypto.DT_BF16) y.axpy_(x, alpha=2.0) # BF16 在内部转换为 FP32 计算

一维场景

y = pypto.tensor([128], pypto.DT_FP32) x = pypto.tensor([128], pypto.DT_FP32) pypto.set_vec_tile_shapes(64) y.axpy_(x, alpha=2.0)

【免费下载链接】pyptoPyPTO（发音: pai p-t-o）：Parallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pypto