Skip to content

Serialize ragged tensors into a single tensor and deserialize them back - #786

Open
10vulture1005 wants to merge 1 commit into
openvinotoolkit:masterfrom
10vulture1005:ragged-tensor-pack-unpack
Open

10vulture1005 wants to merge 1 commit into
openvinotoolkit:masterfrom
10vulture1005:ragged-tensor-pack-unpack

Conversation

@10vulture1005

Copy link
Copy Markdown

Motivation

RaggedTensorPack takes begins, ends and values, but only passed the
values tensor through and ignored the ragged structure entirely. That made it
impossible for a ragged tensor to cross a graph boundary as a single tensor,
which is what the TensorFlow RaggedTensor format does.

What this changes

RaggedTensorPack now serializes all three parts into one 1D u8 tensor, and
a new RaggedTensorUnpack deserializes it back. Together they let a ragged
tensor be a single tensor at a model input or output.

The format is versioned and self-describing:

struct RaggedTensorHeader { // exactly 16 bytes, enforced by static_assert
int32_t version; // format version
int32_t value_type_id; // element::Type_t of values
uint32_t num_elements; // number of rows
uint32_t values_byte_size; // byte size of values
};
int32_t beginsnum_elements;
int32_t endsnum_elements;
uint8_t valuesvalues_byte_size;

It is little-endian independently of the host, and the header is validated on
deserialization, so a truncated or corrupted tensor is reported with a clear
error instead of being misinterpreted. Row offsets are validated against the
values on both the packing and the unpacking side.

pre_translate_ragged_tensor_input() now accepts both the decomposed and the
packed representation, so a frontend translator can consume either.

Two design decisions worth reviewing

1. value_type attribute on RaggedTensorUnpack. The values element type is
stored in the header, i.e. in data only available at runtime, but OpenVINO
requires a static element type on every model output — the CPU plugin rejects
element::dynamic outright. The type is therefore taken from a new value_type
attribute, or from the header when that attribute is absent and the input is a
Constant. evaluate() re-checks the two against each other to catch a
mismatched configuration. Feedback welcome if upstream would prefer a different
mechanism.

2. RaggedTensorPack reports a dynamic output shape. A static lengthcomputed from the input shapes turned out to be wrong in practice: a plugin may
change the precision of an input, and the CPU plugin does convert i64→i32,
i16→i32, f64→f32, u8→u16 for extension ops. The length is now set in
evaluate() from the values actually seen. This trades a little shapeinformation for correctness.

Testing

19 new tests in tests/layer_tests.py:

  • round-trip begins/ends/values → RaggedTensorPack → RaggedTensorUnpack
    over 11 value types (i8/i16/i32/i64/u8/f32/f64/bool/...) and edge cases
    (fully empty, single element, all rows empty, rows of varying length)
  • the packed bytes are compared against an independent hand-built encoding of
    the format, so an accidental layout change fails the test
  • a packed tensor produced by one model is consumed by a separate model
  • 5 corruption cases (truncated header, short body, bad version, unknown type
    id, inconsistent row count) are rejected
  • string values are rejected, since they cannot be serialized into a
    self-contained byte stream

Full tests/layer_tests.py: 159 passed vs. 140 on unmodified HEAD. The 4
failures and 2 errors that remain are pre-existing — I built unmodified HEAD
separately and reproduced them identically (Unicode/ICU data drift, unrelated to
this change).

RaggedTensorPack accepted begins, ends and values but only passed the values
tensor through, so a ragged tensor could not cross a graph boundary as one
tensor the way a TensorFlow RaggedTensor does.

RaggedTensorPack now serializes all three parts into a single 1D u8 tensor and
the new RaggedTensorUnpack deserializes it back, using a versioned,
self-describing little-endian format: a fixed 16-byte header (version, values
element type, number of rows, values byte size) followed by begins, ends and
the raw values. The header is checked on deserialization, so a corrupted or
truncated tensor is reported instead of being misinterpreted, and the row
offsets are validated against the values on both sides.

The values element type is stored in the header, that is in the data that is
only available at runtime, while OpenVINO requires a static element type for
every output of a model. RaggedTensorUnpack therefore takes it either from a
new `value_type` attribute or, when that attribute is absent and the input is a
constant, from the header of that constant, and it verifies that both agree.

The length of the packed tensor is not reported by type inference because it
depends on the values element type as it is seen at execution time, and a
plugin may change the precision of an input, for example by converting i64 to
i32. The length is set by evaluate() instead.

pre_translate_ragged_tensor_input() now accepts both the decomposed and the
packed representation of a ragged tensor, so a translator can consume either.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant