CANN/pypto动态维度API文档

📅 2026/7/4 7:21:50 👁️ 次浏览

pypto.DYNAMIC【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pypto产品支持情况Ascend 950PR支持Atlas A3 训练系列产品/Atlas A3 推理系列产品支持Atlas A2 训练系列产品/Atlas A2 推理系列产品支持功能说明pypto.DYNAMIC用于定义动态维度Dynamic Dimension允许张量的某些维度在运行时变化。这对于处理可变的batch size、序列长度等场景非常有用。动态维度通常在模块级别定义然后在JIT编译的内核函数的类型注解中使用。主要应用场景动态Batch Size: 推理时batch size可能随请求数量变化动态序列长度: NLP任务中文本序列长度不固定动态图结构: 图神经网络中节点数量可变条件计算: 根据输入形状决定计算流程Shape标记方式标记含义pypto.DYNAMIC或pypto.DYN动态轴传入torch tensor该维变化时无需重编译pypto.STATIC静态轴传入torch tensor该维变化时触发重编译64固定轴只允许传入该固定大小传入其他大小会报错(runtime_debug_mode为3开启校验)...剩余轴都作为静态轴处理约束说明动态维度必须在JIT函数的类型注解中使用调用示例示例1: 基础用法 - 动态Batch Sizeimport pypto # 固定轴 HIDDEN_SIZE 128 pypto.frontend.jit def add_bias( x: pypto.Tensor([pypto.DYNAMIC, pypto.STATIC], pypto.DT_FP32), bias: pypto.Tensor([HIDDEN_SIZE], pypto.DT_FP32), out: pypto.Tensor([pypto.DYNAMIC, ...], pypto.DT_FP32) ): # 实现add逻辑 # [pypto.DYNAMIC, ...]第一维是动态的省略号表示剩余维度是静态的 ... # 可以用不同的batch size调用 x1 torch.randn(2, 128, dtypetorch.float32, devicenpu:0) out1 torch.randn(2, 128, dtypetorch.float32, devicenpu:0) result1 add_bias(x1, bias, out1) # batch2 x2 torch.randn(8, 128, dtypetorch.float32, devicenpu:0) out2 torch.randn(8, 128, dtypetorch.float32, devicenpu:0) result2 add_bias(x2, bias, out2) # batch8示例2: 多个动态维度HIDDEN 768 pypto.frontend.jit def attention_kernel( q: pypto.Tensor([pypto.DYNAMIC, pypto.DYNAMIC, HIDDEN], pypto.DT_FP32), k: pypto.Tensor([pypto.DYNAMIC, pypto.DYNAMIC, HIDDEN], pypto.DT_FP32), v: pypto.Tensor([pypto.DYNAMIC, pypto.DYNAMIC, HIDDEN], pypto.DT_FP32), out: pypto.Tensor([pypto.DYNAMIC, pypto.DYNAMIC, HIDDEN], pypto.DT_FP32), ): # 实现attention逻辑 # 前两维batch、序列长度都是动态的 ... return output # 可以处理不同的batch和序列长度 attention_kernel(q_4_128, k_4_128, v_4_128, out) # B4, SEQ128 attention_kernel(q_2_256, k_2_256, v_2_256, out) # B2, SEQ256无需重编译最佳实践文档说明: 在代码注释中说明哪些维度是动态的及其含义测试覆盖: 测试不同的动态维度取值确保代码的正确性【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pypto创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考

相关新闻