[TF FE] Fix sporadic failure in LinSpace (#15570)

Signed-off-by: Kazantsev, Roman <roman.kazantsev@intel.com>
This commit is contained in:
Roman Kazantsev 2023-02-08 12:54:58 +04:00 committed by GitHub
parent d94dae79d8
commit 6582ad7e4d
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -27,10 +27,15 @@ OutputVector translate_linspace_op(const NodeContext& node) {
Output<Node> delta = make_shared<Subtract>(stop, start); Output<Node> delta = make_shared<Subtract>(stop, start);
delta = make_shared<Divide>(delta, num_minus_one); delta = make_shared<Divide>(delta, num_minus_one);
// generate a range of numbers [0, 1, ..., num)
// to have exact numbers of elements equal to num
auto const_zero = make_shared<Constant>(num.get_element_type(), Shape{}, 0);
auto range0_n = make_shared<Range>(const_zero, num, const_one, start.get_element_type());
// compute the result // compute the result
auto stop_plus_delta = make_shared<Add>(stop, delta); Output<Node> linspace = make_shared<Multiply>(range0_n, delta);
auto linspace = make_shared<Range>(start, stop_plus_delta, delta, start.get_element_type()); linspace = make_shared<Add>(linspace, start);
set_node_name(node.get_name(), linspace); set_node_name(node.get_name(), linspace.get_node_shared_ptr());
return {linspace}; return {linspace};
} }
} // namespace op } // namespace op