[TF FE] Refactor Where translator and test it (#13880)

Signed-off-by: Kazantsev, Roman <roman.kazantsev@intel.com>

Signed-off-by: Kazantsev, Roman <roman.kazantsev@intel.com>
This commit is contained in:
Roman Kazantsev
2022-11-07 12:32:16 +04:00
committed by GitHub
parent 2e9a8cd332
commit 995bcdbe47
2 changed files with 27 additions and 101 deletions
+5 -3
View File
@@ -6,6 +6,7 @@
#include "openvino/opsets/opset8.hpp"
using namespace std;
using namespace ov;
using namespace ov::opset8;
namespace ov {
@@ -14,9 +15,10 @@ namespace tensorflow {
namespace op {
OutputVector translate_where_op(const NodeContext& node) {
auto x = node.get_input(0);
auto non_zero = make_shared<NonZero>(x);
auto transpose_order = make_shared<Constant>(element::i64, Shape{2}, vector<int64_t>{1, 0});
default_op_checks(node, 1, {"Where"});
auto condition = node.get_input(0);
auto non_zero = make_shared<NonZero>(condition, element::i64);
auto transpose_order = make_shared<Constant>(element::i32, Shape{2}, vector<int32_t>{1, 0});
auto res = make_shared<opset8::Transpose>(non_zero, transpose_order);
set_node_name(node.get_name(), res);
return res->outputs();