1 // Copyright (C) 2018-2019 Intel Corporation
2 // SPDX-License-Identifier: Apache-2.0
7 #include "ie_built_in_impl.hpp"
14 namespace InferenceEngine {
15 namespace ShapeInfer {
18 *@brief Implementation of Shape inference for TopK layer
20 class TopKShapeProp : public BuiltInShapeInferImpl {
22 explicit TopKShapeProp(const std::string& type) : BuiltInShapeInferImpl(type) {}
24 void inferShapesImpl(const std::vector<Blob::CPtr>& inBlobs,
25 const std::map<std::string, std::string>& params,
26 const std::map<std::string, Blob::Ptr>& blobs,
27 std::vector<SizeVector>& outShapes) override {
29 TopKLayer topKLayer(lp);
30 topKLayer.params = params;
31 topKLayer.type = _type;
32 validate(&topKLayer, inBlobs, params, blobs);
34 const size_t TOPK_DATA = 0;
35 const size_t TOPK_K = 1;
37 if (inBlobs[TOPK_DATA]->getTensorDesc().getPrecision() != Precision::FP32)
38 THROW_IE_EXCEPTION << " Incorrect input data tensor precision. Only FP32 is supported!";
40 if (inBlobs[TOPK_K]->getTensorDesc().getPrecision() != Precision::I32)
41 THROW_IE_EXCEPTION << " Incorrect input index value precision. Only I32 is supported!";
43 if (inBlobs[TOPK_K]->getTensorDesc().getDims().size() > 1)
44 THROW_IE_EXCEPTION << " Index vector should be 1 dimension";
46 SizeVector src_dims = inBlobs[TOPK_DATA]->getTensorDesc().getDims();
47 int axis_ = topKLayer.axis;
49 axis_ += src_dims.size();
51 size_t axis = static_cast<size_t>(axis_);
53 if (src_dims.size() < (1 + axis))
54 THROW_IE_EXCEPTION << " Incorrect input parameters dimensions and axis number!";
56 int *src_k = inBlobs[TOPK_K]->cbuffer().as<int *>();
58 THROW_IE_EXCEPTION << " Only const input for 'k' is supported!";
60 src_k += inBlobs[TOPK_K]->getTensorDesc().getBlockingDesc().getOffsetPadding();
62 outShapes.push_back(inShapes[0]);
63 outShapes.push_back(inShapes[0]);
64 outShapes[0][axis] = static_cast<size_t>(src_k[0]);
65 outShapes[1][axis] = static_cast<size_t>(src_k[0]);
69 } // namespace ShapeInfer
70 } // namespace InferenceEngine