| 23 | namespace xla { |
| 24 | |
| 25 | XlaOp TopK(XlaOp input, int64 k) { |
| 26 | XlaBuilder* const builder = input.builder(); |
| 27 | return builder->ReportErrorOrReturn([&]() -> StatusOr<XlaOp> { |
| 28 | TF_ASSIGN_OR_RETURN(Shape input_shape, builder->GetShape(input)); |
| 29 | int last_dim = input_shape.dimensions_size() - 1; |
| 30 | |
| 31 | Shape iota_shape = |
| 32 | ShapeUtil::MakeShape(S32, AsInt64Slice(input_shape.dimensions())); |
| 33 | XlaOp iota_s32 = Iota(builder, iota_shape, last_dim); |
| 34 | for (int64 i = 0; i < input_shape.rank(); ++i) { |
| 35 | if (input_shape.is_dynamic_dimension(i)) { |
| 36 | // Propagate dynamic dimension from inputs to iota. |
| 37 | iota_s32 = SetDimensionSize(iota_s32, GetDimensionSize(input, i), i); |
| 38 | } |
| 39 | } |
| 40 | auto input_dims = input_shape.dimensions(); |
| 41 | XlaOp sort_result = |
| 42 | Sort({input, iota_s32}, |
| 43 | CreateScalarGtComputation({input_shape.element_type(), S32}, |
| 44 | iota_s32.builder()), |
| 45 | last_dim, /*is_stable=*/true); |
| 46 | std::vector<int64> start_indices(input_shape.dimensions_size(), 0); |
| 47 | std::vector<int64> limit_indices(input_dims.begin(), input_dims.end()); |
| 48 | limit_indices[last_dim] = k; |
| 49 | std::vector<int64> strides(input_shape.dimensions_size(), 1); |
| 50 | |
| 51 | XlaOp values = Slice(GetTupleElement(sort_result, 0), start_indices, |
| 52 | limit_indices, strides); |
| 53 | XlaOp indices = Slice(GetTupleElement(sort_result, 1), start_indices, |
| 54 | limit_indices, strides); |
| 55 | return Tuple(builder, {values, indices}); |
| 56 | }); |
| 57 | } |
| 58 | |
| 59 | XlaOp TopKWithPartitions(XlaOp input, int64 k, int64 num_partitions) { |
| 60 | XlaBuilder* const builder = input.builder(); |