Lines Matching defs:partitions
17 // 1. Let N be the size of partitions. We initialize a new vector indices_in
20 // by partitions and indices_in. This will result in two new vectors
182 // N - the size of the partitions tensor
184 // P - the number of partitions
211 const Tensor* partitions, const Tensor* partition_count,
219 for (int i = partitions->dims(); i < data->dims(); i++) {
229 const Tensor& partitions = c->input(1);
232 c, TensorShapeUtils::StartsWith(data.shape(), partitions.shape()),
234 "data.shape must start with partitions.shape, ",
236 ", partitions.shape = ", partitions.shape().DebugString()),
241 // We must handle the case of empty partitions separately,
243 if (partitions.NumElements() == 0) {
254 this->AllocateOutputs(c, &data, &partitions, &partition_count, &outputs,
268 // Also sort the info in partitions and output it in indices_out,
270 this->CountAndSortParts(c, &partitions, &partition_count, &indices_out,
303 auto wrapped_callback = [this, c, &data, &partitions, indices_out,
306 this->AllocateOutputs(c, &data, &partitions, &cpu_tensor, &outputs, done);
311 int32 N = partitions.NumElements();
323 void RadixSort(OpKernelContext* c, const Tensor* partitions,
326 int32 N = partitions->NumElements();
333 const int32* partitions_ptr = partitions->flat<int32>().data();
357 void CountAndSortParts(OpKernelContext* c, const Tensor* partitions,
362 int32 N = partitions->NumElements();
371 this->RadixSort(c, partitions, &indices_in, &partitions_out, indices_out,
375 // each index appears in partitions.
421 // each index appears in partitions. The distinct indices are stored
429 // at least once in partitions. We move each value from aggregates_out
