Home
last modified time | relevance | path

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

/external/tensorflow/tensorflow/compiler/xla/client/lib/
Dsorting_test.cc83 xla::GetTupleElement(xla::TopKWithPartitions(x, 3, /*num_partitions=*/2), 0); in XLA_TEST_F()
91 xla::GetTupleElement(xla::TopKWithPartitions(x_rev, 3, /*num_partitions=*/2), in XLA_TEST_F()
100 xla::GetTupleElement(xla::TopKWithPartitions(x, 3, /*num_partitions=*/3), 0); in XLA_TEST_F()
108 xla::GetTupleElement(xla::TopKWithPartitions(x_rev, 3, /*num_partitions=*/3), in XLA_TEST_F()
117 xla::GetTupleElement(xla::TopKWithPartitions(x, 3, /*num_partitions=*/5), 0); in XLA_TEST_F()
138 xla::GetTupleElement(xla::TopKWithPartitions(x_rev, 3, /*num_partitions=*/5), in XLA_TEST_F()
147 xla::GetTupleElement(xla::TopKWithPartitions(a, 3, /*num_partitions=*/2), 1); in XLA_TEST_F()
Dsorting.h31 XlaOp TopKWithPartitions(XlaOp input, int64 k, int64 num_partitions = 1);
Dsorting.cc44 return TopKWithPartitions(input, k, num_partitions); in TopK()
76 XlaOp TopKWithPartitions(XlaOp input, int64 k, int64 num_partitions) { in TopKWithPartitions() function