
PyPTO 张量散射写入详解pypto.Tensor.scatter 的用法、参数与实现原理【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pypto导读pypto.Tensor.scatter是 PyPTOParallel Tensor/Tile Operation 编程范式中用于按索引将源数据写入目标张量的核心接口与gather按索引取值互为逆操作广泛用于索引更新、稀疏赋值、one-hot 打散等场景。本文基于 docs/zh/api/tensor_api/tensor/pypto-Tensor-scatter.md 及其关联的 pypto.scatter_ 操作文档完整覆盖函数原型、参数约束、reduce 归约模式、TileShape 切分规则与可运行示例并辅以仓库源码印证其底层调用链帮助读者在 Atlas 训练/推理系列产品上正确、高效地使用 scatter 接口。产品支持情况pypto.Tensor.scatter含非 inplace 版本pypto.scatter与 inplace 版本pypto.scatter_在当前仓库支持以下产品Ascend 950PR / Ascend 950DT支持Atlas A3 训练系列产品 / Atlas A3 推理系列产品支持Atlas A2 训练系列产品 / Atlas A2 推理系列产品支持功能说明scatter将src的值写入input张量写入位置由index张量指定。以 3 维张量为例计算公式如下其他维度以此类推src为固定标量时$$ \begin{cases} input\left[ index\left[i\right]\left[j\right]\left[k\right] \right]\left[j\right]\left[k\right] src \text{if } dim 0 \ input\left[i\right]\left[ index\left[i\right]\left[j\right]\left[k\right] \right]\left[k\right] src \text{if } dim 1 \ input\left[i\right]\left[j\right]\left[ index\left[i\right]\left[j\right]\left[k\right] \right] src \text{if } dim 2 \end{cases} $$src为 Tensor 时$$ \begin{cases} input\left[ index\left[i\right]\left[j\right]\left[k\right] \right]\left[j\right]\left[k\right] src\left[i\right]\left[j\right]\left[k\right] \text{if } dim 0 \ input\left[i\right]\left[ index\left[i\right]\left[j\right]\left[k\right] \right]\left[k\right] src\left[i\right]\left[j\right]\left[k\right] \text{if } dim 1 \ input\left[i\right]\left[j\right]\left[ index\left[i\right]\left[j\right]\left[k\right] \right] src\left[i\right]\left[j\right]\left[k\right] \text{if } dim 2 \end{cases} $$从公式可见dim决定index中的索引值作用于input的哪一个轴其余轴保持对齐。在 python/pypto/op/indexing.py 的 docstring 中这一语义被表述为self[index[i][j][k]][j][k] srcdim0等三种形式两者完全一致。函数原型pypto.Tensor.scatter是张量成员方法原型为scatter(self, dim: int, index: Tensor, src: float, *, reduce: str None) - Tensor对应的函数形式pypto.scatter原型为scatter(input: Tensor, dim: int, index: Tensor, src: Union[float, Element, Tensor], *, reduce: str None) - Tensor其中scatter是scatter_的non-inplace非原地版本scatter_直接在input上修改并返回更新后的input原地操作而scatter返回一个新张量。两个接口在 python/pypto/tensor.py 中均通过source_location装饰器直接转发到pypto.scatter_/pypto.scatter。参数说明参数名输入/输出说明input输入支持类型Tensor。支持的数据类型DT_FP32、DT_FP16、DT_BF16、DT_INT8、DT_UINT8、DT_INT16、DT_INT32、DT_INT64。不支持空 TensorShape 仅支持 1~4 维Shape Size 不大于 2147483647即 INT32_MAX。dim输入指定用于索引的维度支持 input 维度范围内的任意维度。合法的维度索引范围为-input.dim到input.dim - 1支持负索引。index输入input 的一组索引。支持类型Tensor数据类型为 INT64 或 INT32。维度需与 input 保持一致对所有的d ! dim维度需满足index.size(d) input.size(d)当 src 为 Tensor 时所有维度都需满足index.size(d) src.size(d)。不支持空 TensorShape Size 不大于 2147483647。src输入更新的标量或 Tensor。src 为 Element标量时支持的数据类型为 DT_FP32、DT_FP16、DT_BF16、DT_INT8、DT_UINT8、DT_INT16、DT_INT32、DT_INT64不支持输入 INF/NANsrc 为 Tensor 时支持的数据类型同上且数据类型需与 input 保持一致。reduce输入要应用的归约操作支持add或multiply不传参None时默认为直接替换。其中index的数据类型与src的类型校验在 python/pypto/op/indexing.py 中实现index必须是 INT32/INT64否则抛出TypeError错误码 0xF00001src必须是 int/float、Element 或 Tensor否则同样抛出类型错误。reduce 归约模式说明reduce参数控制写入时的归约行为仓库源码 python/pypto/op/indexing.py 中的get_scatter_mode函数将其映射为底层 ScatterMode 枚举reduceNone直接替换ScatterMode.NONE默认行为reduceadd累加写入ScatterMode.ADD即input[index] srcreducemultiply乘累写入ScatterMode.MULTIPLY即input[index] * src其他字符串抛出PyptoError错误码 0xF00002ValueError(scatter reduce only support add, multiply)。返回值说明scatter_返回更新后的input原地操作返回同一对象底层通过input.Move(...)完成就地更新。scatter返回一个包含 scatter 结果的新张量不修改原input。约束说明broadcast 约束input 和 index 不支持 broadcastviewshape 约束input.shape 的 dim 轴不可切viewshape 的维度与 input 维度相同要求viewshape[dim] max(input.shape[dim], index.shape[dim])其余维度的 Shape 大小不做限制tileshape 约束input.shape 的 dim 轴不可切tileshape 的维度与 input 维度相同tileshape[dim] viewshape[dim]其余维度的 Shape 大小不做限制。input、index 和 result 都会放在 UB统一缓冲区中需满足所有输入和输出的 tileshape 大小总和不能超过 UB 内存的大小切分对齐input.shape 和 index.shape 的非 dim 轴切分需满足 viewshape[non dim] 切分后input 和 index 的非 dim 轴切分块数相同tileshape 切分时也需要保证 input 和 index 的非 dim 轴切分块数相同不唯一索引src 为 Tensor、reduce 为 None且 index 中存在指向相同位置的不唯一索引时行为是不确定的将从 src 中任意选择一个值写入。约束 2、3 中“dim 轴不可切”的要求在底层实现 framework/src/interface/operation/vector/scatter.cpp 中有直接体现切分时校验vecTile[axis] dstTensor-shape[axis]与vecTile[axis] idxInput-shape[axis]若违反则报错ERR_CONFIG_TILEThe axis is not supported for tile splitting并将 dim 轴的 tile 大小强制设为max(dstShape[axis], idxShape[axis])以保证全载。TileShape 设置示例调用 scatter 接口前应通过set_vec_tile_shapes设置 TileShapeTileShape 维度应与输出一致。例如输入 input shape 为[a, b, c]dim 为 1index 为[m, t, p]其中 mapcsrc 为[x, y, z]其中 xmytzp输出为[a, b, c]TileShape 设置为[m1, t1, p1]。则 m1、p1 分别用于切分 m、p 轴t1 必须大于等于 b 和 tdim 对应轴不可切必须保证 b 轴和 t 轴全载pypto.set_vec_tile_shapes(4, 16, 32)调用示例将 2 维 input 根据 2 维 index 更新对应索引的值x pypto.tensor([3, 5], pypto.DT_FP32) y pypto.tensor([2, 2], pypto.DT_INT64) o pypto.scatter_(x, 0, y, 2.0)结果示例如下输入数据x:[[0 0 0 0 0], [0 0 0 0 0], [0 0 0 0 0]] 输入数据y:[[1 2], [0 1]] 输出数据o:[[2.0 0 0 0 0], [2.0 2.0 0 0 0], [0 2.0 0 0 0]]解读dim0 表示index中的数值作用于第 0 维。y[0][0]1、y[0][1]2、y[1][0]0、y[1][1]1于是o[1][0]、o[2][1]、o[0][0]、o[1][1]被写入标量 2.0其余位置保持 0。该示例同样被记录在 python/pypto/op/indexing.py 的 docstring 中可直接作为单元级验证用例参考。其他常见用法张量方法形式inplacex.scatter_(dim, index, src, reduceadd)见 python/pypto/tensor.py张量方法形式non-inplacex.scatter(dim, index, src)见 python/pypto/tensor.pyTensor 作为 srcpypto.scatter_(x, dim, index, src_tensor)要求 src 与 input 数据类型一致归约模式pypto.scatter_(x, 0, y, 1.0, reduceadd)实现按索引累加。底层实现与调用链从源码结构可以梳理出 scatter 的完整调用链Python 层python/pypto/op/indexing.py 中的scatter_/scatter完成类型校验与 reduce 模式解析后调用pypto_impl.Scatter(input, index, src, dim, scatter_mode)绑定层pypto_impl.Scatter为 C 绑定接口见 python/src/bindings/ 相关实现将参数透传给框架层框架层framework/src/interface/operation/vector/scatter.cpp 中InnerTiledScatterElementS按 TileShape 递归切分 dstTensor对每个 tile 调用function.AddOperation(Opcode::OP_SCATTER_ELEMENT, {srcTile, idxTile}, {dstTile})并写入axis、scalar、scatter_mode三个属性最终由后端代码生成模块vector 相关 codegen如 framework/src/codegen/npu/codegen_vector_gather_scatter.cpp翻译为昇腾 NPU 可执行的向量指令。值得注意的是Tensor 的scatter_update接口python/pypto/op/indexing.py对应另一个算子族framework/src/interface/operation/vector/scatter_update.cpp其语义是“将 src 中 index 指定的元素更新到 input”与本文的 scatter按 index 将 src 写入 input不同使用时需注意区分。与 gather 的关系gather是 scatter 的逆操作gather 沿dim指定的轴按index从 input 中取出值而 scatter 按index将 src 写入 input。在 python/pypto/op/indexing.py 的 docstring 中scatter 明确标注了See Also: gather: The inverse operation, gather values along an axis specified by dim.。二者配合可实现完整的“取-写”索引数据流是 PyPTO 中张量索引类操作的基础能力之一。总结pypto.Tensor.scatter及函数形式pypto.scatter/ inplace 形式pypto.scatter_为 PyPTO 提供了按索引散射写入张量的能力支持标量、Element 与 Tensor 三类 src支持add/multiply两种归约模式可运行于 Ascend 950 系列与 Atlas A2/A3 系列产品。使用时需重点把握三点dim 轴不可切分、非 dim 轴切分块数需 input 与 index 对齐、UB 内存需容纳全部 tile 输入与输出同时注意 reduce 为 None 且 index 含重复索引时结果不确定。结合 pypto.scatter_ 的参数表格与本文调用示例即可在算子开发中正确落地该接口。【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pypto创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考