Skip to content

[WiP] Extend Support for RaggedTensorToTensor #613

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
apaniukov wants to merge 3 commits into
base: master
Choose a base branch
from
+202 −56
Open

[WiP] Extend Support for RaggedTensorToTensor #613

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
9eff335
WiP
apaniukov Feb 20, 2026
735b792
Fix shape infer
p-wysocki Mar 6, 2026
a0b223f
Merge pull request #2 from p-wysocki/ragged_tsr
apaniukov Mar 6, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
138 changes: 92 additions & 46 deletions src/ragged_to_dense.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,8 +14,8 @@ void RaggedToDense::validate_and_infer_types() {
OPENVINO_ASSERT(get_input_size() == 3 + 1 + 1 || get_input_size() == 3 + 1 + 1 + 1,
"RaggedToDense requires 5 inputs (begins, ends, data, padding_size, value) and 1 optional input (pad_right).");

// Input ragged tensor
check_ragged_input(this, 0);
// Input ragged tensor (begins, ends, data)
check_ragged_input_any_rank_data(this, 0);

// Target size along ragged dimension
OPENVINO_ASSERT(get_input_element_type(3).is_integral_number());
Expand All @@ -33,18 +33,29 @@ void RaggedToDense::validate_and_infer_types() {

set_input_is_relevant_to_shape(3);

if(get_input_partial_shape(0).rank().is_dynamic()) {
const auto begins_shape = get_input_partial_shape(0);
const auto data_shape = get_input_partial_shape(2);
const auto begins_rank = begins_shape.rank();
const auto data_rank = data_shape.rank();

if (begins_rank.is_dynamic() || data_rank.is_dynamic()) {
set_output_type(0, get_input_element_type(2), PartialShape::dynamic());
set_output_type(1, element::boolean, PartialShape::dynamic());
} else {
auto shape = get_input_partial_shape(0);
if(auto target_dim = dynamic_cast<Constant*>(get_input_node_ptr(3))) {
shape.push_back(target_dim->cast_vector<int64_t>()[0]);
auto out_shape = begins_shape;
if (auto target_dim = dynamic_cast<Constant*>(get_input_node_ptr(3))) {
out_shape.push_back(target_dim->cast_vector<int64_t>()[0]);
} else {
shape.push_back(Dimension());
out_shape.push_back(Dimension());
}

const auto data_rank_len = static_cast<size_t>(data_rank.get_length());
for (size_t idx = 1; idx < data_rank_len; ++idx) {
out_shape.push_back(data_shape[idx]);
}
set_output_type(0, get_input_element_type(2), shape);
set_output_type(1, element::boolean, shape);

set_output_type(0, get_input_element_type(2), out_shape);
set_output_type(1, element::boolean, out_shape);
}
if (get_input_size() == 3 + 1 + 1 + 1) {
OPENVINO_ASSERT(get_input_partial_shape(5).is_dynamic() || get_input_partial_shape(5).is_static() && get_input_partial_shape(5).rank().get_length() == 0,
Expand All @@ -60,18 +71,40 @@ bool RaggedToDense::evaluate(ov::TensorVector& outputs, const ov::TensorVector&
// FIXME: Works for POD types only (not for strings!)
// FIXME: Output mask is calculated even if there are no consumers
auto begins = inputs[0].data<const int32_t>();
auto ends = inputs[1].data<const int32_t>();
auto nelems = inputs[0].get_size();
auto elems = reinterpret_cast<const char*>(inputs[2].data());
auto elem_size = inputs[2].get_element_type().size();
auto default_value = reinterpret_cast<const char*>(inputs[4].data());
auto ends = inputs[1].data<const int32_t>();
const auto nelems = inputs[0].get_size();

const auto elems = reinterpret_cast<const char*>(inputs[2].data());
const auto elem_size = inputs[2].get_element_type().size();
const auto default_value = reinterpret_cast<const char*>(inputs[4].data());

// Suppose validate was called and set correct output shape
// Take a target shape value for ragged dimension
size_t target_dim = inputs[3].data<const int32_t>()[0];
const size_t target_dim = static_cast<size_t>(inputs[3].data<const int32_t>()[0]);

// If output shape is dynamic at compile-time (e.g. target_dim is not constant),
// set it at runtime based on actual input values.
{
ov::Shape out_shape = inputs[0].get_shape();
out_shape.push_back(target_dim);
const auto& data_shape = inputs[2].get_shape();
for (size_t idx = 1; idx < data_shape.size(); ++idx) {
out_shape.push_back(data_shape[idx]);
}
outputs[0].set_shape(out_shape);
outputs[1].set_shape(out_shape);
}

// Number of dense elements per one ragged element (trailing dense dimensions).
// For 1D data, this equals 1.
const auto& data_shape = inputs[2].get_shape();
size_t inner_elems = 1;
for (size_t idx = 1; idx < data_shape.size(); ++idx) {
inner_elems *= data_shape[idx];
}

auto out_elems = reinterpret_cast<char*>(outputs[0].data());
auto out_mask = outputs[1].data<char>();
// Mask may be unallocated when this op is used as an intermediate node with no consumers
auto out_mask = outputs[1] ? outputs[1].data<char>() : nullptr;

auto out_elem_orig = out_elems;
auto out_mask_orig = out_mask;
Expand All @@ -81,48 +114,61 @@ bool RaggedToDense::evaluate(ov::TensorVector& outputs, const ov::TensorVector&
pad_right = inputs[5].data<bool>()[0];
}

auto fill_default_block = [&](char*& dst) {
for (size_t k = 0; k < inner_elems; ++k) {
dst = std::copy(default_value, default_value + elem_size, dst);
}
};

auto mask_fill = [&](char*& dst, size_t count, char value) {
if (dst) {
dst = std::fill_n(dst, count, value);
}
};

if (pad_right) {
for(size_t i = 0; i < nelems; ++i) {
auto begin = elems + elem_size * begins[i];
auto target_len = (
std::min(size_t(ends[i] - begins[i]), target_dim) * (1 - m_pad_max_length) // truncation
+ target_dim * m_pad_max_length // pad to max length
);
auto end = begin + elem_size * target_len;
for (size_t i = 0; i < nelems; ++i) {
const size_t data_len = static_cast<size_t>(ends[i] - begins[i]);
size_t target_len = (std::min(data_len, target_dim) * (1 - m_pad_max_length) +
target_dim * m_pad_max_length);

const auto begin = elems + elem_size * inner_elems * static_cast<size_t>(begins[i]);
const auto end = begin + elem_size * inner_elems * target_len;
out_elems = std::copy(begin, end, out_elems);
out_mask = std::fill_n(out_mask, target_len, char(1));
if(target_len < target_dim)
out_mask = std::fill_n(out_mask, target_dim - target_len, char(0));
while(target_len < target_dim) {
out_elems = std::copy(default_value, default_value + elem_size, out_elems);

mask_fill(out_mask, target_len * inner_elems, char(1));
if (target_len < target_dim) {
mask_fill(out_mask, (target_dim - target_len) * inner_elems, char(0));
}

while (target_len < target_dim) {
fill_default_block(out_elems);
++target_len;
}
}
} else {
for(size_t i = 0; i < nelems; ++i) {
const size_t data_len = ends[i] - begins[i];
auto target_len = (
std::min(data_len, target_dim) * (1 - m_pad_max_length) // truncation
+ target_dim * m_pad_max_length // pad to max length
);
auto pad_len = target_dim - target_len;

// fill padding values
for (size_t i = 0; i < nelems; ++i) {
const size_t data_len = static_cast<size_t>(ends[i] - begins[i]);
size_t target_len = (std::min(data_len, target_dim) * (1 - m_pad_max_length) +
target_dim * m_pad_max_length);
const size_t pad_len = target_dim - target_len;

for (size_t j = 0; j < pad_len; ++j) {
out_elems = std::copy(default_value, default_value + elem_size, out_elems);
fill_default_block(out_elems);
}
// fill actual values
auto begin = elems + elem_size * begins[i];
auto end = begin + elem_size * target_len;

const auto begin = elems + elem_size * inner_elems * static_cast<size_t>(begins[i]);
const auto end = begin + elem_size * inner_elems * target_len;
out_elems = std::copy(begin, end, out_elems);

// construct padding mask
out_mask = std::fill_n(out_mask, pad_len, char(0));
out_mask = std::fill_n(out_mask, target_dim - pad_len, char(1));
mask_fill(out_mask, pad_len * inner_elems, char(0));
mask_fill(out_mask, (target_dim - pad_len) * inner_elems, char(1));
}
}

OPENVINO_ASSERT(out_elems == out_elem_orig + outputs[0].get_byte_size());
OPENVINO_ASSERT(out_mask == out_mask_orig + outputs[1].get_byte_size());
if (out_mask) {
OPENVINO_ASSERT(out_mask == out_mask_orig + outputs[1].get_byte_size());
}
return true;
}
20 changes: 18 additions & 2 deletions src/ragged_to_ragged.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,8 +19,24 @@ void RaggedToRagged::validate_and_infer_types() {
OPENVINO_ASSERT(rowids_type == element::i32, "Expected an i32 rowids tensor ragged representation.");
OPENVINO_ASSERT(first_dim_size_type == element::i32, "Expected an i32 first dim size tensor ragged representation.");

set_output_type(0, get_input_element_type(0), PartialShape({ Dimension::dynamic() }));
set_output_type(1, get_input_element_type(0), PartialShape({ Dimension::dynamic() }));
// Check whether input 1 is a Constant node, otherwise fall back to the lower-bound tensor
PartialShape out_shape({ Dimension::dynamic() });
auto infer_from_int32_tensor = [&](const ov::Tensor& t) {
if (t && t.get_element_type() == element::i32 && t.get_size() >= 1) {
out_shape = PartialShape({ static_cast<Dimension::value_type>(t.data<const int32_t>()[0]) });
}
};
if (const auto* first_dim_const = dynamic_cast<const Constant*>(get_input_node_ptr(1))) {
auto vals = first_dim_const->cast_vector<int32_t>();
if (!vals.empty()) {
out_shape = PartialShape({ static_cast<Dimension::value_type>(vals[0]) });
}
} else {
infer_from_int32_tensor(get_input_tensor(1).get_lower_value());
}

set_output_type(0, get_input_element_type(0), out_shape);
set_output_type(1, get_input_element_type(0), out_shape);
}


Expand Down
8 changes: 8 additions & 0 deletions src/ragged_to_ragged.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,4 +38,12 @@ class RaggedToRagged : public ov::op::Op {
bool has_evaluate() const override {
return true;
}

bool evaluate_lower(ov::TensorVector& output_values) const override {
return false;
}

bool evaluate_upper(ov::TensorVector& output_values) const override {
return false;
}
};
80 changes: 72 additions & 8 deletions src/tensorflow_translators.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -392,17 +392,23 @@ ov::OutputVector translate_ragged_tensor_to_tensor(const ov::frontend::NodeConte
auto default_value = node.get_input(2);
auto row_partition_types = node.get_attribute<std::vector<std::string>>("row_partition_types");

FRONT_END_GENERAL_CHECK((row_partition_types == std::vector<std::string>{"ROW_SPLITS"} && node_input_size == 4) ||
(row_partition_types == std::vector<std::string>{"FIRST_DIM_SIZE", "VALUE_ROWIDS"} && node_input_size == 5),
"[TensorFlow Frontend] internal error: RaggedTensorToTensor is supported only for 2D tensor "
"with ROW_SPLITS or {FIRST_DIM_SIZE, VALUE_ROWIDS} type");
FRONT_END_GENERAL_CHECK(
(row_partition_types == std::vector<std::string>{"ROW_SPLITS"} && node_input_size == 4) ||
(row_partition_types == std::vector<std::string>{"FIRST_DIM_SIZE", "VALUE_ROWIDS"} && node_input_size == 5) ||
(row_partition_types == std::vector<std::string>{"ROW_SPLITS", "VALUE_ROWIDS"} && node_input_size == 5),
"[TensorFlow Frontend] internal error: RaggedTensorToTensor is supported only for 2D tensor "
"with ROW_SPLITS or {FIRST_DIM_SIZE, VALUE_ROWIDS} type, or for 3D tensor with {ROW_SPLITS, VALUE_ROWIDS} type");

// shape can be undefined (with -1 value) or defined with positive values
// in this case replace_shape will be selected below

// since begins, ends and target shape are expected to be of int32 type
shape = std::make_shared<Convert>(shape, ov::element::i32);

// ensure pad value is a scalar
auto scalar_shape = std::make_shared<Constant>(ov::element::i32, Shape{ 0 }, std::vector<int32_t>{});
default_value = std::make_shared<Reshape>(default_value, scalar_shape, false);

ov::Output<ov::Node> begins, ends;
ov::Output<ov::Node> longest_batch, longest_row_size;
if (row_partition_types == std::vector<std::string>{"ROW_SPLITS"}) {
Expand All @@ -427,7 +433,7 @@ ov::OutputVector translate_ragged_tensor_to_tensor(const ov::frontend::NodeConte
auto reduce_axis = std::make_shared<Constant>(ov::element::i32, Shape{ 1 }, 0);
longest_row_size = std::make_shared<ReduceMax>(longest_row_size, reduce_axis, true);
}
else {
else if (row_partition_types == std::vector<std::string>{"FIRST_DIM_SIZE", "VALUE_ROWIDS"}) {
auto first_dim_size = node.get_input(3);
auto value_rowids = node.get_input(4);

Expand Down Expand Up @@ -459,6 +465,66 @@ ov::OutputVector translate_ragged_tensor_to_tensor(const ov::frontend::NodeConte
longest_row_size = std::make_shared<ReduceMax>(longest_row_size, reduce_axis, true);
}

if (row_partition_types == std::vector<std::string>{"ROW_SPLITS", "VALUE_ROWIDS"}) {
// Two ragged dimensions:
// - outer: ROW_SPLITS
// - inner: VALUE_ROWIDS
auto row_splits = node.get_input(3);
row_splits = std::make_shared<Convert>(row_splits, ov::element::i32);

// Outer begins/ends and outer batch size
auto row_splits_shape = std::make_shared<ShapeOf>(row_splits, ov::element::i32)->output(0);
auto const_one = std::make_shared<Constant>(ov::element::i32, Shape{}, 1);
auto outer_batch = std::make_shared<Subtract>(row_splits_shape, const_one)->output(0);
auto begins_start = std::make_shared<Constant>(ov::element::i32, Shape{ 1 }, 0);
auto ends_start = std::make_shared<Constant>(ov::element::i32, Shape{ 1 }, 1);
auto step = std::make_shared<Constant>(ov::element::i32, Shape{ 1 }, 1);
auto outer_begins = std::make_shared<Slice>(row_splits, begins_start, outer_batch, step);
auto outer_ends = std::make_shared<Slice>(row_splits, ends_start, row_splits_shape, step);

auto reduce_axis0 = std::make_shared<Constant>(ov::element::i32, Shape{ 1 }, 0);
auto outer_max_len = std::make_shared<ReduceMax>(
std::make_shared<Subtract>(outer_ends, outer_begins)->output(0), reduce_axis0, true);

// Inner rows total = row_splits[last]
auto gather_axis0 = std::make_shared<Constant>(ov::element::i32, Shape{ 1 }, 0);
auto last_index = std::make_shared<Subtract>(row_splits_shape, const_one)->output(0);
auto inner_rows_total = std::make_shared<Gather>(row_splits, last_index, gather_axis0)->output(0);

auto value_rowids = node.get_input(4);
value_rowids = std::make_shared<Convert>(value_rowids, ov::element::i32);

// Build begins/ends for inner rows from value_rowids
auto inner_r2r = std::make_shared<RaggedToRagged>(ov::OutputVector{ value_rowids, inner_rows_total });
auto inner_begins = inner_r2r->output(0);
auto inner_ends = inner_r2r->output(1);
auto inner_max_len = std::make_shared<ReduceMax>(
std::make_shared<Subtract>(inner_ends, inner_begins)->output(0), reduce_axis0, true);

// First densify inner ragged dimension -> [inner_rows_total, inner_max_len]
auto inner_dense =
std::make_shared<RaggedToDense>(ov::OutputVector{ inner_begins, inner_ends, values, inner_max_len, default_value })->output(0);

// Then densify outer ragged dimension over the inner-dense tensor
auto outer_dense =
std::make_shared<RaggedToDense>(ov::OutputVector{ outer_begins, outer_ends, inner_dense, outer_max_len, default_value })->output(0);

auto replace_shape_3d =
std::make_shared<Concat>(ov::OutputVector{ outer_batch, outer_max_len, inner_max_len }, 0)->output(0);
auto const_zero = std::make_shared<Constant>(ov::element::i32, Shape{}, 0);
auto shape_less_zero = std::make_shared<Less>(shape, const_zero);
shape = std::make_shared<Select>(shape_less_zero, replace_shape_3d, shape);

auto pads_begin = std::make_shared<Constant>(ov::element::i32, Shape{ 3 }, std::vector<int32_t>{0, 0, 0});
auto pads_end = std::make_shared<Subtract>(shape, replace_shape_3d);
auto result_dense_tensor =
std::make_shared<Pad>(outer_dense, pads_begin, pads_end, default_value, ov::op::PadMode::CONSTANT)->output(0);

result_dense_tensor.get_node_shared_ptr()->set_friendly_name(node_name);
result_dense_tensor.set_names({ node_name + ":0" });
return { result_dense_tensor };
}

auto ragged_to_dense = std::make_shared<RaggedToDense>(ov::OutputVector{ begins, ends, values, longest_row_size, default_value })->output(0);

// adjust shape value since it can contain -1 value that means a dimension must be deduced based on minimal dimension size
Expand All @@ -472,9 +538,7 @@ ov::OutputVector translate_ragged_tensor_to_tensor(const ov::frontend::NodeConte
// note that replace_shape to be equal a shape of ragged_to_dense
// Pad operation removes (or crops) if padding number is negative
auto pads_end = std::make_shared<Subtract>(shape, replace_shape);
auto squeeze_axis = std::make_shared<Constant>(ov::element::i32, Shape{ 1 }, 0);
auto pad_value = std::make_shared<Squeeze>(default_value, squeeze_axis);
auto result_dense_tensor = std::make_shared<Pad>(ragged_to_dense, pads_begin, pads_end, pad_value, ov::op::PadMode::CONSTANT)->output(0);
auto result_dense_tensor = std::make_shared<Pad>(ragged_to_dense, pads_begin, pads_end, default_value, ov::op::PadMode::CONSTANT)->output(0);

result_dense_tensor.get_node_shared_ptr()->set_friendly_name(node_name);
result_dense_tensor.set_names({ node_name + ":0" });
Expand Down
10 changes: 10 additions & 0 deletions src/utils.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -62,6 +62,16 @@ void check_ragged_input(const Node* node, size_t input_index) {
FRONT_END_GENERAL_CHECK(rank.is_dynamic() || rank.get_length() == 1, "The last tensor in ragged tensor representation should be a 1D tensor");
}

void check_ragged_input_any_rank_data(const Node* node, size_t input_index) {
FRONT_END_GENERAL_CHECK(node->get_input_element_type(input_index + 0) == element::i32,
"Expected an i32 tensor as the first part of the decomposed ragged representation");
FRONT_END_GENERAL_CHECK(node->get_input_element_type(input_index + 1) == element::i32,
"Expected an i32 tensor as the second part of the decomposed ragged representation");
auto rank = node->get_input_partial_shape(input_index + 2).rank();
FRONT_END_GENERAL_CHECK(rank.is_dynamic() || rank.get_length() >= 1,
"The last tensor in ragged tensor representation should have rank >= 1");
}

void check_ragged_string_input(const Node* node, size_t input_index) {
FRONT_END_GENERAL_CHECK(node->get_input_element_type(input_index+0) == element::i32, "Expected an i32 tensor as the first part of the decomposed ragged string representation");
FRONT_END_GENERAL_CHECK(node->get_input_element_type(input_index+1) == element::i32, "Expected an i32 tensor as the second part of the decomposed ragged string representation");
Expand Down
2 changes: 2 additions & 0 deletions src/utils.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,8 @@ void check_string_scalar_input(const ov::Node* node, size_t input_index);

void check_ragged_input(const ov::Node* node, size_t input_index);

void check_ragged_input_any_rank_data(const ov::Node* node, size_t input_index);

void check_ragged_string_input(const ov::Node* node, size_t input_index);

void set_string_output(ov::Node* node, size_t output_index, const ov::PartialShape& shape);
Expand Down
Loading