#include "../node_context.h"
#include "../op_table.h"
#include "../utils.h"
#include <openvino/op/constant.hpp>
#include <openvino/op/greater.hpp>
#include <openvino/op/greater_eq.hpp>
#include <openvino/op/less.hpp>
#include <openvino/op/less_eq.hpp>
#include <openvino/op/range.hpp>
#include <openvino/op/reshape.hpp>
#include <openvino/op/select.hpp>
namespace ov {
namespace frontend {
namespace ggml {
namespace op {
OutputVector translate_tri(const NodeContext & context) {
num_inputs_check(context, 1, 1);
auto x = context.get_input(0);
int32_t tri_type = context.get_output_op_params()[0];
auto shape = context.get_input_shape(0).to_shape();
int64_t n = static_cast<int64_t>(shape[3]);
auto start = ov::op::v0::Constant::create(ov::element::i64, {}, {int64_t(0)});
auto stop = ov::op::v0::Constant::create(ov::element::i64, {}, {n});
auto step = ov::op::v0::Constant::create(ov::element::i64, {}, {int64_t(1)});
auto range = std::make_shared<ov::op::v4::Range>(start, stop, step, ov::element::i64);
auto col_shape = ov::op::v0::Constant::create(ov::element::i64, {4}, std::vector<int64_t>{1, 1, 1, n});
auto col_idx = std::make_shared<ov::op::v1::Reshape>(range, col_shape, false);
auto row_shape = ov::op::v0::Constant::create(ov::element::i64, {4}, std::vector<int64_t>{1, 1, n, 1});
auto row_idx = std::make_shared<ov::op::v1::Reshape>(range, row_shape, false);
std::shared_ptr<ov::Node> mask;
switch (tri_type) {
case 0: mask = std::make_shared<ov::op::v1::GreaterEqual>(col_idx, row_idx);
break;
case 1: mask = std::make_shared<ov::op::v1::Greater>(col_idx, row_idx);
break;
case 2: mask = std::make_shared<ov::op::v1::LessEqual>(col_idx, row_idx);
break;
case 3: mask = std::make_shared<ov::op::v1::Less>(col_idx, row_idx);
break;
default:
throw std::runtime_error("translate_tri: invalid tri_type " + std::to_string(tri_type));
}
auto zero = ov::op::v0::Constant::create(ov::element::f32, {}, {0.0f});
auto res = std::make_shared<ov::op::v1::Select>(mask, x, zero);
return rename_outputs_with_suffix({res}, context.get_name());
}
} } } }