| 990 | //===----------------------------------------------------------------------===// |
| 991 | |
| 992 | static void BuildTopKOp(Builder *builder, OperationState &result, Value input, |
| 993 | Value k) { |
| 994 | // Output size is only known if k is constant value. A negative dimension is |
| 995 | // considered dynamic so use -1 here if k is not a constant value. |
| 996 | int const_k = -1; |
| 997 | ElementsAttr cst; |
| 998 | if (matchPattern(k, m_Constant(&cst))) |
| 999 | // These casts should all be valid due to how Tensor constants are stored. |
| 1000 | // TODO(jpienaar): This should use a helper function. |
| 1001 | const_k = cst.getValue<IntegerAttr>({}).getValue().getSExtValue(); |
| 1002 | |
| 1003 | auto val_type = input.getType().cast<TensorType>(); |
| 1004 | // If value is unranked, then so is results. |
| 1005 | if (!val_type.hasRank()) |
| 1006 | return TFL::TopKV2Op::build( |
| 1007 | builder, result, UnrankedTensorType::get(val_type.getElementType()), |
| 1008 | UnrankedTensorType::get(builder->getIntegerType(32)), input, k); |
| 1009 | |
| 1010 | // Resultant shape is value.shape[:-1] + [k] |
| 1011 | std::vector<int64_t> shape(val_type.getShape()); |
| 1012 | shape[shape.size() - 1] = const_k; |
| 1013 | TFL::TopKV2Op::build( |
| 1014 | builder, result, RankedTensorType::get(shape, val_type.getElementType()), |
| 1015 | RankedTensorType::get(shape, builder->getIntegerType(32)), input, k); |
| 1016 | } |
| 1017 | |
| 1018 | //===----------------------------------------------------------------------===// |
| 1019 | // FakeQuantOp |