CANN/pypto条件操作API

pypto.cond【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pypto产品支持情况产品是否支持Atlas A3 训练系列产品/Atlas A3 推理系列产品√Atlas A2 训练系列产品/Atlas A2 推理系列产品√功能说明定义一个if 条件操作实现python中的if功能。函数原型cond(scalar: SymInt)参数说明参数名输入/输出说明scalar输入条件表达式可以是整数或SymbolicScalar符号标量用于判断条件是否为真返回值说明pypto_impl.RecordIfBranch: 返回一个条件分支对象用于 Python 的 if 语句约束说明必须与 Python 的 if、elif、else 语句配合使用条件表达式会被记录到计算图中支持嵌套的条件语句当函数未使用 pypto.frontend.jit 或 pypto.frontend.function 装饰器修饰时条件表达式需要用 pypto.cond 包装调用示例# 未使用装饰器需要用 pypto.cond 包装条件表达式 def kernel(): ... for s2_idx in pypto.loop(0, 10, 1, power_of_2(max_unroll_times), nameLOOP_L0_bIdx_mla_prolog, idx_nameb_idx): if pypto.cond(pypto.is_loop_end(s2_idx, bn_per_batch)): ... # 使用装饰器无需 pypto.cond 包装 pypto.frontend.jit def kernel(): ... for s2_idx in pypto.loop(0, 10, 1, power_of_2(max_unroll_times), nameLOOP_L0_bIdx_mla_prolog, idx_nameb_idx): if pypto.is_loop_end(s2_idx, bn_per_batch): ...【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pypto创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考