SplitToSequence op with string tensor inputs behaves incorrectly in ONNX Runtime 1.17.1

[Craigacp]

Describe the issue

The SplitToSequence operator behaves incorrectly in ORT 1.17.1 when keepdims=0 and running on string sequences (not sure if both are necessary, but we have other uses of SplitToSequence which don't fail which don't set keepdims and aren't on strings). Instead of splitting the tensor along the specified axis of size n and returning a sequence of length n it just returns the first element n times. This op works correctly in ORT 1.16.3.

As far as I can tell this is the only PR which touches the relevant op between v1.16.3 and v1.17.1 - #18594.

cc @yuslepukhin

To reproduce

Running the following code in ORT 1.17 produces an incorrect output:

import numpy as np
import onnxruntime

test_input = np.array(["This is a test String.", "this is a different and much longer string"])

sess = onnxruntime.InferenceSession(
    "split-to-sequence.onnx", 
    providers=onnxruntime.get_available_providers())
out = sess.run(["split_inputs"], {"input_strings": test_input})
out

gives the output:

[[array('This is a test String.', dtype=object),
  array('This is a test String.', dtype=object)]]

whereas on ORT 1.16.3 it gives:

[[array('This is a test String.', dtype=object),
  array('this is a different and much longer string', dtype=object)]]

Rename split-to-sequence.onnx.txt to split-to-sequence.onnx.
split-to-sequence.onnx.txt

The ONNX model was generated with this code:

from onnx import TensorProto, helper as onnx_helper
input_values = [
    onnx_helper.make_tensor_value_info("input_strings", elem_type=TensorProto.STRING, shape=["batch_size"])]
split = onnx_helper.make_node("SplitToSequence", inputs=["input_strings"], outputs=["split_inputs"], keepdims=0,
                              name="split_inputs")
output_values = [onnx_helper.make_value_info("split_inputs", type_proto=onnx_helper.make_sequence_type_proto(onnx_helper.make_tensor_type_proto(elem_type=TensorProto.STRING,shape=["A"])))]
graph = onnx_helper.make_graph(nodes=[split], name="split_to_sequence", inputs=input_values,
                               outputs=output_values)

extension_opset = onnx_helper.make_operatorsetid("com.microsoft.extensions", 1000)
contrib_opset = onnx_helper.make_operatorsetid("ai.onnx.contrib", 1)
default_opset = onnx_helper.make_operatorsetid("", 17)
opset = [default_opset, contrib_opset, extension_opset]

model = onnx_helper.make_model(graph, opset_imports=opset, ir_version=8)
with open("split-to-sequence.onnx", "wb") as f:
    f.write(model.SerializeToString())

Urgency

Yes, this blocks batch use of ONNX Runtime when embedding strings using BERT tokenizers in our system and batches are important for building the vector index quickly.

Platform

Linux

OS Version

Oracle Linux 9

ONNX Runtime Installation

Released Package

ONNX Runtime Version or Commit ID

v1.17.1

ONNX Runtime API

Python

Architecture

X64

Execution Provider

Default CPU

Execution Provider Library Version

No response

[sophies927]

@pranavsharma can you look into this?

[Craigacp]

After a bit more testing it doesn't seem related to the keepdims argument, I think it's specific to strings. Looks like there's a special case for string inputs in the code, I'll try to put together a test for that op with strings.

[Craigacp]

This test case which I added to onnxruntime/test/providers/cpu/sequence/sequence_ops_test.cc triggers the bug:

TEST(SequenceOpsTest, SplitToSequence_PositiveAxisStringSplit) {
  OpTester test("SplitToSequence", 11);
  test.AddInput<std::string>("input", {2}, std::vector<std::string>({"Test string", "Another string"}));
  int64_t axis = 0;
  test.AddAttribute("axis", axis);
  SeqTensors<std::string> output;
  output.AddTensor({1}, {"Test string"});
  output.AddTensor({1}, {"Another string"});
  test.AddSeqOutput("S2", output);
  test.Run();
}

Error message:

onnxruntime/onnxruntime/test/providers/checkers.cc:71: Failure
Expected equality of these values:
  cur_expected[i]
    Which is: "Another string"
  cur_actual[i]
    Which is: "Test string"
i:0
Google Test trace:
onnxruntime/onnxruntime/test/providers/checkers.cc:484: provider type: CPUExecutionProvider
onnxruntime/onnxruntime/test/providers/base_tester.cc:791: registered execution providers: CPUExecutionProvider
Stack trace:
  0x10c0a92b7: onnxruntime::utils::MLTypeCallDispatcher<>::InvokeWithLeadingTemplateArgs<>()
  0x10c0a3d8a: onnxruntime::test::Check<>()
  0x10c0a4033: onnxruntime::test::CheckOrtValuesAreEqual()
  0x10c09db8f: onnxruntime::test::BaseTester::ExecuteModel<>()
  0x10c09c43f: onnxruntime::test::BaseTester::ExecuteModelForEps()
  0x10c09a9c6: onnxruntime::test::BaseTester::RunWithConfig()
  0x10c099c56: onnxruntime::test::BaseTester::Run()
  0x10c4e676c: onnxruntime::test::SequenceOpsTest_SplitToSequence_PositiveAxisStringSplit_Test::TestBody()
  0x10d386719: testing::internal::HandleExceptionsInMethodIfSupported<>()
  0x10d386632: testing::Test::Run()
  0x10d387c20: testing::TestInfo::Run()
... Google Test internal frames ...

[  FAILED  ] SequenceOpsTest.SplitToSequence_PositiveAxisStringSplit (3 ms)

An otherwise identical test which splits a single dimensional float tensor passes, so it's something in the string special case, but unfortunately my C++ knowledge isn't good enough to figure out exactly what it's doing.