DynamicPartition
【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库，实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math
产品支持情况
| 产品 | 是否支持 |
|---|---|
| Ascend 950PR/Ascend 950DT | √ |
| Atlas A3 训练系列产品/Atlas A3 推理系列产品 | × |
| Atlas A2 训练系列产品/Atlas A2 推理系列产品 | × |
| Atlas 200I/500 A2 推理产品 | × |
| Atlas 推理系列产品 | × |
| Atlas 训练系列产品 | × |
功能说明
接口功能：用于根据分区索引将输入数据动态分割成多个张量。示例： 假设输入Tensor为[1, 6, 3, 8, 2, 9]，将一维张量按条件分成两组（如元素值≤5和>5），值如下所示：
>>> x = tf.constant([1, 6, 3, 8, 2, 9]) >>> partitions = tf.cast(data ≤ 5, tf.int32) >>> tf.dynamic_partition(x, partitions, num_partitions=2) [array([1, 3, 2]), array([6, 8, 9])]支持对多维张量的分区操作，例如将矩阵按行或列拆分：
>>> x = tf.reshape(tf.range(12), [3, 4]) // 按行分区 >>> tf.dynamic_partition(matrix, [0, 1, 0], num_partitions=2) array([[0, 1, 2, 3], [8, 9, 10, 11]])参数说明
| 参数名 | 输入/输出/属性 | 描述 | 数据类型 | 数据格式 |
|---|---|---|---|---|
| x | 输入 | 待分区的原始输入张量 | FLOAT、DOUBLE、FLOAT16、COMPLEX64、UINT8、INT8、INT16、INT32、INT64、UINT16、UINT32、UINT64、BOOL、BFLOAT16、HIFLOAT8、FLOAT8_E5M2、FLOAT8_E4M3FN | ND |
| partitions | 输入 | 分区索引张量，决定每个元素归属的分区编号 | INT32 | ND |
| num_partitions | 可选属性 |
| INT64 | ND |
| y | 输出 | 分区后的输出张量 | FLOAT、DOUBLE、FLOAT16、COMPLEX64、UINT8、INT8、INT16、INT32、INT64、UINT16、UINT32、UINT64、BOOL、BFLOAT16、HIFLOAT8、FLOAT8_E5M2、FLOAT8_E4M3FN | ND |
约束说明
无
调用示例
| 调用方式 | 样例代码 | 说明 |
|---|---|---|
| 图模式调用 | test_geir_dynamic_partition | 通过算子IR构图方式调用DynamicPartition算子。 |
【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库，实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math
创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
