Searched refs:dim_partition_counts (Results 1 – 2 of 2) sorted by relevance
209 auto dim_partition_counts = ShapePartitionAssigner(instruction->shape()) in AssignParallelTasksHelper() local212 ShapePartitionAssigner::GetTotalPartitionCount(dim_partition_counts); in AssignParallelTasksHelper()225 new_root->set_outer_dimension_partitions(dim_partition_counts); in AssignParallelTasksHelper()
170 std::vector<int64> dim_partition_counts(num_outer_dims_to_partition); in TEST_F() local179 dim_partition_counts[i] = dim_partition_count; in TEST_F()184 ShapePartitionIterator partition_iterator(shape, dim_partition_counts); in TEST_F()