Home
last modified time | relevance | path

Searched refs:per_partition_size (Results 1 – 3 of 3) sorted by relevance

/external/tensorflow/tensorflow/compiler/xla/client/lib/
Dsorting.cc84 const int64 per_partition_size = CeilOfRatio(last_dim_size, num_partitions); in TopKWithPartitions() local
86 if (k >= per_partition_size) { in TopKWithPartitions()
110 ConstantR0<int32>(builder, per_partition_size)); in TopKWithPartitions()
112 DynamicSliceInMinorDims(input, {start}, {per_partition_size}); in TopKWithPartitions()
114 DynamicSliceInMinorDims(iota_s32, {start}, {per_partition_size}); in TopKWithPartitions()
145 limit_indices[last_dim] = per_partition_size; in TopKWithPartitions()
/external/tensorflow/tensorflow/compiler/xla/service/spmd/
Dspmd_partitioner.cc1665 const int64 per_partition_size = CeilOfRatio(input_size, partition_count); in HandleSort() local
1683 replicated_dimensions[sort_dim] = per_partition_size * partition_count; in HandleSort()
1825 const int64 per_partition_size = CeilOfRatio(input_size, shard_count); in HandleCustomCall() local
1827 if (k >= per_partition_size) { in HandleCustomCall()
1882 LiteralUtil::CreateR0<int32>(per_partition_size))))), in HandleCustomCall()
2384 const int64 per_partition_size = partitioned_shape.dimensions(dim); in HandleDynamicUpdateSlice() local
2387 if ((partitioned_slice_offsets[i] / per_partition_size) != in HandleDynamicUpdateSlice()
2389 per_partition_size)) { in HandleDynamicUpdateSlice()
2398 LiteralUtil::CreateR0<int>(per_partition_size))); in HandleDynamicUpdateSlice()
Dspmd_partitioner_util.cc1251 const int64 per_partition_size = CeilOfRatio(input_size, shard_count); in GetKValueInTopKWhenPartitionSortDim() local
1253 if (k.value() >= per_partition_size) { in GetKValueInTopKWhenPartitionSortDim()