DynamicPartition

产品支持情况

产品	是否支持
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	可选属性	指定总分区的数量 默认值为1。	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算子。