// Copyright (c) ONNX Project Contributors
//
// SPDX-License-Identifier: Apache-2.0

#include "onnx_core/backend_test/test_case.h"
#include "onnx_extensions/backend_test/cases/sequence/include_sequence_cases.h"

#include <gtest/gtest.h>

#include <cstdint>
#include <string>
#include <vector>

using namespace ONNX_LIGHT_NAMESPACE;
using onnx_backend_test::CollectSequenceTestCases;

namespace {
std::vector<core::backend_test::TestCase> CollectTestCases(const std::string &op_type = "") {
  std::vector<core::backend_test::TestCase> registry;
  CollectSequenceTestCases(registry, op_type);
  return registry;
}
} // namespace
using core::backend_test::TestCase;

namespace Test {

TEST(BackendTestCase, SequenceConstructCaseIsPresent) {
  auto cases = CollectTestCases("SequenceConstruct");
  const TestCase *seq_case = nullptr;
  for (const auto &c : cases) {
    if (c.name == "test_cc_sequence_construct") {
      seq_case = &c;
    }
  }
  ASSERT_NE(seq_case, nullptr);

  const GraphProto &graph = seq_case->model().ref_graph();
  ASSERT_EQ(graph.ref_node().size(), 1u);
  const NodeProto &node = graph.ref_node()[0];
  const auto &op_type = node.ref_op_type();
  EXPECT_EQ(std::string(op_type.data(), op_type.size()), "SequenceConstruct");
  EXPECT_EQ(graph.ref_input().size(), 3u);
  ASSERT_EQ(graph.ref_output().size(), 1u);

  // Output value-info must be promoted to Sequence<Tensor<FLOAT, [2, 3]>>.
  const ValueInfoProto &out_vi = graph.ref_output()[0];
  const TypeProto &out_tp = out_vi.ref_type();
  ASSERT_TRUE(out_tp.has_sequence_type());
  EXPECT_FALSE(out_tp.has_tensor_type());
  const TypeProto &elem = out_tp.ref_sequence_type().ref_elem_type();
  ASSERT_TRUE(elem.has_tensor_type());
  EXPECT_EQ(elem.ref_tensor_type().ref_elem_type(), core::runtime::DataType::FLOAT);

  ASSERT_EQ(seq_case->data_sets().size(), 1u);
  const auto &ds = seq_case->data_sets()[0];
  ASSERT_EQ(ds.inputs.size(), 3u);
  ASSERT_EQ(ds.outputs.size(), 1u);
  EXPECT_EQ(ds.outputs[0].data_type, static_cast<int32_t>(core::runtime::DataType::FLOAT));
  const std::vector<int64_t> expected_shape = {3, 2, 3};
  EXPECT_EQ(ds.outputs[0].shape, expected_shape);
  // Stacked output bytes equal the concatenation of the per-input buffers.
  std::vector<uint8_t> expected_bytes;
  for (const auto &in : ds.inputs) {
    expected_bytes.insert(expected_bytes.end(), in.data.begin(), in.data.end());
  }
  EXPECT_EQ(ds.outputs[0].data, expected_bytes);
}

TEST(BackendTestCase, ConcatFromSequenceCasesAreRegistered) {
  auto cases = CollectTestCases();
  const TestCase *axis_0 = nullptr;
  const TestCase *axis_1 = nullptr;
  const TestCase *new_axis = nullptr;
  for (const auto &c : cases) {
    if (c.name == "test_cc_concat_from_sequence_axis_0") {
      axis_0 = &c;
    } else if (c.name == "test_cc_concat_from_sequence_axis_1") {
      axis_1 = &c;
    } else if (c.name == "test_cc_concat_from_sequence_new_axis") {
      new_axis = &c;
    }
  }
  ASSERT_NE(axis_0, nullptr);
  ASSERT_NE(axis_1, nullptr);
  ASSERT_NE(new_axis, nullptr);

  // All three cases share the same in-graph topology: a
  // SequenceConstruct that bundles tensor inputs into a sequence, then
  // a ConcatFromSequence node that consumes the sequence.
  for (const TestCase *tc : {axis_0, axis_1, new_axis}) {
    const GraphProto &graph = tc->model().ref_graph();
    ASSERT_EQ(graph.ref_node().size(), 2u) << tc->name;
    const NodeProto &seq_node = graph.ref_node()[0];
    EXPECT_EQ(seq_node.ref_op_type(), "SequenceConstruct");
    EXPECT_EQ(seq_node.ref_input().size(), 3u);
    EXPECT_EQ(seq_node.ref_output().size(), 1u);

    const NodeProto &concat_node = graph.ref_node()[1];
    EXPECT_EQ(concat_node.ref_op_type(), "ConcatFromSequence");
    ASSERT_EQ(concat_node.ref_input().size(), 1u);
    EXPECT_EQ(concat_node.ref_input()[0], seq_node.ref_output()[0]);
    ASSERT_EQ(concat_node.ref_output().size(), 1u);
    EXPECT_EQ(concat_node.ref_output()[0], "concat_result");

    ASSERT_EQ(graph.ref_input().size(), 3u);
    ASSERT_EQ(graph.ref_output().size(), 1u);
    const ValueInfoProto &out_vi = graph.ref_output()[0];
    EXPECT_EQ(out_vi.ref_type().has_tensor_type(), true);

    ASSERT_EQ(tc->data_sets().size(), 1u);
    const auto &ds = tc->data_sets()[0];
    ASSERT_EQ(ds.inputs.size(), 3u);
    ASSERT_EQ(ds.outputs.size(), 1u);
    EXPECT_EQ(ds.outputs[0].data_type, static_cast<int32_t>(core::runtime::DataType::FLOAT));
  }

  // Expected output shapes per case.
  EXPECT_EQ(axis_0->data_sets()[0].outputs[0].shape, (std::vector<int64_t>{6, 3}));
  EXPECT_EQ(axis_1->data_sets()[0].outputs[0].shape, (std::vector<int64_t>{2, 9}));
  EXPECT_EQ(new_axis->data_sets()[0].outputs[0].shape, (std::vector<int64_t>{3, 2, 3}));
}

TEST(BackendTestCase, SequenceLengthCaseIsPresent) {
  auto cases = CollectTestCases("SequenceLength");
  const TestCase *length_case = nullptr;
  for (const auto &c : cases) {
    if (c.name == "test_cc_sequence_length") {
      length_case = &c;
      break;
    }
  }
  ASSERT_NE(length_case, nullptr);

  const GraphProto &graph = length_case->model().ref_graph();
  ASSERT_EQ(graph.ref_node().size(), 2u);
  EXPECT_EQ(std::string(graph.ref_node()[0].ref_op_type()), "SequenceConstruct");
  EXPECT_EQ(std::string(graph.ref_node()[1].ref_op_type()), "SequenceLength");
  ASSERT_EQ(graph.ref_output().size(), 1u);
  const ValueInfoProto &out_vi = graph.ref_output()[0];
  ASSERT_TRUE(out_vi.ref_type().has_tensor_type());
  EXPECT_EQ(out_vi.ref_type().ref_tensor_type().ref_elem_type(), core::runtime::DataType::INT64);
  EXPECT_EQ(out_vi.ref_type().ref_tensor_type().ref_shape().ref_dim().size(), 0u);

  ASSERT_EQ(length_case->data_sets().size(), 1u);
  const auto &ds = length_case->data_sets()[0];
  ASSERT_EQ(ds.inputs.size(), 3u);
  ASSERT_EQ(ds.outputs.size(), 1u);
  EXPECT_EQ(ds.outputs[0].data_type, static_cast<int32_t>(core::runtime::DataType::INT64));
  EXPECT_TRUE(ds.outputs[0].shape.empty());
  ASSERT_EQ(ds.outputs[0].data.size(), sizeof(int64_t));
  EXPECT_EQ(*ds.outputs[0].AsInt64(), 3);
}

TEST(BackendTestCase, SequenceEraseCasesAreRegistered) {
  auto cases = CollectTestCases("SequenceErase");
  const TestCase *default_case = nullptr;
  const TestCase *pos1_case = nullptr;
  const TestCase *neg_case = nullptr;
  for (const auto &c : cases) {
    if (c.name == "test_cc_sequence_erase_default") {
      default_case = &c;
    } else if (c.name == "test_cc_sequence_erase_pos1") {
      pos1_case = &c;
    } else if (c.name == "test_cc_sequence_erase_neg") {
      neg_case = &c;
    }
  }
  ASSERT_NE(default_case, nullptr);
  ASSERT_NE(pos1_case, nullptr);
  ASSERT_NE(neg_case, nullptr);

  // All cases share the same in-graph topology: SequenceConstruct → SequenceErase.
  for (const TestCase *tc : {default_case, pos1_case, neg_case}) {
    const GraphProto &graph = tc->model().ref_graph();
    ASSERT_GE(graph.ref_node().size(), 2u) << tc->name;
    EXPECT_EQ(std::string(graph.ref_node()[0].ref_op_type()), "SequenceConstruct");
    EXPECT_EQ(std::string(graph.ref_node()[1].ref_op_type()), "SequenceErase");
    ASSERT_EQ(graph.ref_output().size(), 1u);
    // Output is a sequence type.
    EXPECT_TRUE(graph.ref_output()[0].ref_type().has_sequence_type()) << tc->name;
    ASSERT_EQ(tc->data_sets().size(), 1u);
    // Output stacked tensor has 2 elements: shape [2, 2, 3].
    EXPECT_EQ(tc->data_sets()[0].outputs[0].shape, (std::vector<int64_t>{2, 2, 3})) << tc->name;
  }

  // The default case has 3 tensor inputs (no position).
  EXPECT_EQ(default_case->data_sets()[0].inputs.size(), 3u);
  // The position cases have 4 inputs (3 tensors + 1 position scalar).
  EXPECT_EQ(pos1_case->data_sets()[0].inputs.size(), 4u);
  EXPECT_EQ(neg_case->data_sets()[0].inputs.size(), 4u);
}

TEST(BackendTestCase, SequenceInsertCasesAreRegistered) {
  auto cases = CollectTestCases("SequenceInsert");
  const TestCase *default_case = nullptr;
  const TestCase *pos1_case = nullptr;
  const TestCase *neg_case = nullptr;
  for (const auto &c : cases) {
    if (c.name == "test_cc_sequence_insert_default") {
      default_case = &c;
    } else if (c.name == "test_cc_sequence_insert_pos1") {
      pos1_case = &c;
    } else if (c.name == "test_cc_sequence_insert_neg") {
      neg_case = &c;
    }
  }
  ASSERT_NE(default_case, nullptr);
  ASSERT_NE(pos1_case, nullptr);
  ASSERT_NE(neg_case, nullptr);

  for (const TestCase *tc : {default_case, pos1_case, neg_case}) {
    const GraphProto &graph = tc->model().ref_graph();
    ASSERT_GE(graph.ref_node().size(), 2u) << tc->name;
    EXPECT_EQ(std::string(graph.ref_node()[0].ref_op_type()), "SequenceConstruct");
    EXPECT_EQ(std::string(graph.ref_node()[1].ref_op_type()), "SequenceInsert");
    ASSERT_EQ(graph.ref_output().size(), 1u);
    EXPECT_TRUE(graph.ref_output()[0].ref_type().has_sequence_type()) << tc->name;
    ASSERT_EQ(tc->data_sets().size(), 1u);
    EXPECT_EQ(tc->data_sets()[0].outputs[0].shape, (std::vector<int64_t>{4, 2, 3})) << tc->name;
  }

  EXPECT_EQ(default_case->data_sets()[0].inputs.size(), 4u);
  EXPECT_EQ(pos1_case->data_sets()[0].inputs.size(), 5u);
  EXPECT_EQ(neg_case->data_sets()[0].inputs.size(), 5u);
}

TEST(BackendTestCase, SequenceAtCasesAreRegistered) {
  auto cases = CollectTestCases("SequenceAt");
  const TestCase *pos0_case = nullptr;
  const TestCase *pos2_case = nullptr;
  const TestCase *neg_case = nullptr;
  for (const auto &c : cases) {
    if (c.name == "test_cc_sequence_at_pos0") {
      pos0_case = &c;
    } else if (c.name == "test_cc_sequence_at_pos2") {
      pos2_case = &c;
    } else if (c.name == "test_cc_sequence_at_neg") {
      neg_case = &c;
    }
  }
  ASSERT_NE(pos0_case, nullptr);
  ASSERT_NE(pos2_case, nullptr);
  ASSERT_NE(neg_case, nullptr);

  // All cases share the same in-graph topology: SequenceConstruct → SequenceAt.
  for (const TestCase *tc : {pos0_case, pos2_case, neg_case}) {
    const GraphProto &graph = tc->model().ref_graph();
    ASSERT_GE(graph.ref_node().size(), 2u) << tc->name;
    EXPECT_EQ(std::string(graph.ref_node()[0].ref_op_type()), "SequenceConstruct");
    EXPECT_EQ(std::string(graph.ref_node()[1].ref_op_type()), "SequenceAt");
    ASSERT_EQ(graph.ref_output().size(), 1u);
    // Output is a tensor type (not a sequence).
    EXPECT_TRUE(graph.ref_output()[0].ref_type().has_tensor_type()) << tc->name;
    ASSERT_EQ(tc->data_sets().size(), 1u);
    // Each selected element has shape [2, 3].
    EXPECT_EQ(tc->data_sets()[0].outputs[0].shape, (std::vector<int64_t>{2, 3})) << tc->name;
    // Inputs are 3 tensors + 1 position scalar.
    EXPECT_EQ(tc->data_sets()[0].inputs.size(), 4u) << tc->name;
  }
}

TEST(BackendTestCase, SequenceMapCasesAreRegistered) {
  auto cases = CollectTestCases("SequenceMap");
  const TestCase *float_case = nullptr;
  const TestCase *int64_case = nullptr;
  const TestCase *identity_2_seq = nullptr;
  const TestCase *add_2_seq = nullptr;
  const TestCase *add_1_seq_1_tensor = nullptr;
  const TestCase *extract_shapes = nullptr;
  for (const auto &c : cases) {
    if (c.name == "test_cc_sequence_map_identity_float") {
      float_case = &c;
    } else if (c.name == "test_cc_sequence_map_identity_int64") {
      int64_case = &c;
    } else if (c.name == "test_cc_sequence_map_identity_2_sequences") {
      identity_2_seq = &c;
    } else if (c.name == "test_cc_sequence_map_add_2_sequences") {
      add_2_seq = &c;
    } else if (c.name == "test_cc_sequence_map_add_1_sequence_1_tensor") {
      add_1_seq_1_tensor = &c;
    } else if (c.name == "test_cc_sequence_map_extract_shapes") {
      extract_shapes = &c;
    }
  }
  ASSERT_NE(float_case, nullptr);
  ASSERT_NE(int64_case, nullptr);
  ASSERT_NE(identity_2_seq, nullptr);
  ASSERT_NE(add_2_seq, nullptr);
  ASSERT_NE(add_1_seq_1_tensor, nullptr);
  ASSERT_NE(extract_shapes, nullptr);

  for (const TestCase *tc : {float_case, int64_case}) {
    const GraphProto &graph = tc->model().ref_graph();
    ASSERT_GE(graph.ref_node().size(), 2u) << tc->name;
    EXPECT_EQ(std::string(graph.ref_node()[0].ref_op_type()), "SequenceConstruct");
    EXPECT_EQ(std::string(graph.ref_node()[1].ref_op_type()), "SequenceMap");

    // The SequenceMap node carries a 'body' GRAPH attribute.
    const NodeProto &map_node = graph.ref_node()[1];
    ASSERT_EQ(map_node.ref_attribute().size(), 1u) << tc->name;
    EXPECT_EQ(std::string(map_node.ref_attribute()[0].ref_name()), "body");
    EXPECT_EQ(map_node.ref_attribute()[0].ref_type(), AttributeProto::AttributeType::GRAPH);

    ASSERT_EQ(graph.ref_output().size(), 1u);
    // Output is a sequence type.
    EXPECT_TRUE(graph.ref_output()[0].ref_type().has_sequence_type()) << tc->name;
    ASSERT_EQ(tc->data_sets().size(), 1u);
  }

  // The FLOAT case has 3 input tensors of shape [2, 3] → stacked [3, 2, 3].
  EXPECT_EQ(float_case->data_sets()[0].inputs.size(), 3u);
  EXPECT_EQ(float_case->data_sets()[0].outputs[0].shape, (std::vector<int64_t>{3, 2, 3}));
  // The INT64 case has 2 input tensors of shape [4] → stacked [2, 4].
  EXPECT_EQ(int64_case->data_sets()[0].inputs.size(), 2u);
  EXPECT_EQ(int64_case->data_sets()[0].outputs[0].shape, (std::vector<int64_t>{2, 4}));

  // identity_2_sequences: 2 SequenceConstruct + 1 SequenceMap → 2 output sequences.
  {
    const GraphProto &g = identity_2_seq->model().ref_graph();
    ASSERT_EQ(g.ref_node().size(), 3u);
    EXPECT_EQ(std::string(g.ref_node()[0].ref_op_type()), "SequenceConstruct");
    EXPECT_EQ(std::string(g.ref_node()[1].ref_op_type()), "SequenceConstruct");
    EXPECT_EQ(std::string(g.ref_node()[2].ref_op_type()), "SequenceMap");
    ASSERT_EQ(g.ref_output().size(), 2u);
    EXPECT_TRUE(g.ref_output()[0].ref_type().has_sequence_type());
    EXPECT_TRUE(g.ref_output()[1].ref_type().has_sequence_type());
    // 6 graph inputs (3 per sequence), 2 stacked outputs.
    ASSERT_EQ(identity_2_seq->data_sets().size(), 1u);
    EXPECT_EQ(identity_2_seq->data_sets()[0].inputs.size(), 6u);
    ASSERT_EQ(identity_2_seq->data_sets()[0].outputs.size(), 2u);
    EXPECT_EQ(identity_2_seq->data_sets()[0].outputs[0].shape, (std::vector<int64_t>{3, 3}));
    EXPECT_EQ(identity_2_seq->data_sets()[0].outputs[1].shape, (std::vector<int64_t>{3, 4}));
  }

  // add_2_sequences: 2 SequenceConstruct + 1 SequenceMap → 1 output sequence.
  {
    const GraphProto &g = add_2_seq->model().ref_graph();
    ASSERT_EQ(g.ref_node().size(), 3u);
    EXPECT_EQ(std::string(g.ref_node()[2].ref_op_type()), "SequenceMap");
    const NodeProto &map_node = g.ref_node()[2];
    // Body is Add (2 inputs, 1 output).
    const GraphProto &body = map_node.ref_attribute()[0].ref_g();
    ASSERT_EQ(body.ref_node().size(), 1u);
    EXPECT_EQ(std::string(body.ref_node()[0].ref_op_type()), "Add");
    ASSERT_EQ(g.ref_output().size(), 1u);
    EXPECT_TRUE(g.ref_output()[0].ref_type().has_sequence_type());
    EXPECT_EQ(add_2_seq->data_sets()[0].inputs.size(), 6u);
    EXPECT_EQ(add_2_seq->data_sets()[0].outputs[0].shape, (std::vector<int64_t>{3, 4}));
    // Verify element-wise add for the first iteration.
    const auto &out = add_2_seq->data_sets()[0].outputs[0];
    const float *out_data = out.As<float>();
    EXPECT_FLOAT_EQ(out_data[0], 10.0f);
    EXPECT_FLOAT_EQ(out_data[3], 13.0f);
  }

  // add_1_sequence_1_tensor: 1 SequenceConstruct + 1 SequenceMap, with a
  // broadcast tensor as the second SequenceMap input.
  {
    const GraphProto &g = add_1_seq_1_tensor->model().ref_graph();
    ASSERT_EQ(g.ref_node().size(), 2u);
    EXPECT_EQ(std::string(g.ref_node()[0].ref_op_type()), "SequenceConstruct");
    EXPECT_EQ(std::string(g.ref_node()[1].ref_op_type()), "SequenceMap");
    const NodeProto &map_node = g.ref_node()[1];
    ASSERT_EQ(map_node.ref_input().size(), 2u);
    // 3 sequence-element tensors + 1 broadcast tensor → 4 graph inputs.
    EXPECT_EQ(add_1_seq_1_tensor->data_sets()[0].inputs.size(), 4u);
    EXPECT_EQ(add_1_seq_1_tensor->data_sets()[0].outputs[0].shape, (std::vector<int64_t>{3, 4}));
    const auto &out = add_1_seq_1_tensor->data_sets()[0].outputs[0];
    const float *out_data = out.As<float>();
    // x0_0 = [0,1,2,3], x1 = [100,200,300,400] → y0_0 = [100,201,302,403].
    EXPECT_FLOAT_EQ(out_data[0], 100.0f);
    // y0_0[3] = x0_0[3] + x1[3] = 3 + 400 = 403.
    EXPECT_FLOAT_EQ(out_data[3], 403.0f);
  }

  // extract_shapes: 1 SequenceConstruct + 1 SequenceMap with Shape body
  // and 3 inputs of distinct rank-3 shapes → INT64[3, 3] output.
  {
    const GraphProto &g = extract_shapes->model().ref_graph();
    ASSERT_EQ(g.ref_node().size(), 2u);
    EXPECT_EQ(std::string(g.ref_node()[1].ref_op_type()), "SequenceMap");
    const NodeProto &map_node = g.ref_node()[1];
    const GraphProto &body = map_node.ref_attribute()[0].ref_g();
    ASSERT_EQ(body.ref_node().size(), 1u);
    EXPECT_EQ(std::string(body.ref_node()[0].ref_op_type()), "Shape");
    ASSERT_EQ(g.ref_output().size(), 1u);
    EXPECT_TRUE(g.ref_output()[0].ref_type().has_sequence_type());
    // Output sequence element type must be INT64.
    EXPECT_EQ(g.ref_output()[0]
                  .ref_type()
                  .ref_sequence_type()
                  .ref_elem_type()
                  .ref_tensor_type()
                  .ref_elem_type(),
              core::runtime::DataType::INT64);
    EXPECT_EQ(extract_shapes->data_sets()[0].inputs.size(), 3u);
    EXPECT_EQ(extract_shapes->data_sets()[0].outputs[0].shape, (std::vector<int64_t>{3, 3}));
    const auto &out = extract_shapes->data_sets()[0].outputs[0];
    EXPECT_EQ(out.data_type, static_cast<int32_t>(core::runtime::DataType::INT64));
    const int64_t *out_data = out.As<int64_t>();
    // Shapes are {4,3,2}, {2,5,1}, {1,1,7}.
    EXPECT_EQ(out_data[0], 4);
    EXPECT_EQ(out_data[1], 3);
    EXPECT_EQ(out_data[2], 2);
    EXPECT_EQ(out_data[3], 2);
    EXPECT_EQ(out_data[4], 5);
    EXPECT_EQ(out_data[5], 1);
    EXPECT_EQ(out_data[6], 1);
    EXPECT_EQ(out_data[7], 1);
    EXPECT_EQ(out_data[8], 7);
  }
}

TEST(BackendTestCase, SequenceEmptyCasesAreRegistered) {
  auto cases = CollectTestCases("SequenceEmpty");
  const TestCase *default_case = nullptr;
  const TestCase *int64_case = nullptr;
  for (const auto &c : cases) {
    if (c.name == "test_cc_sequence_empty_default") {
      default_case = &c;
    } else if (c.name == "test_cc_sequence_empty_int64") {
      int64_case = &c;
    }
  }
  ASSERT_NE(default_case, nullptr);
  ASSERT_NE(int64_case, nullptr);

  // Both cases share the same in-graph topology: SequenceEmpty → SequenceLength.
  for (const TestCase *tc : {default_case, int64_case}) {
    const GraphProto &graph = tc->model().ref_graph();
    ASSERT_EQ(graph.ref_node().size(), 2u) << tc->name;
    EXPECT_EQ(std::string(graph.ref_node()[0].ref_op_type()), "SequenceEmpty");
    EXPECT_EQ(std::string(graph.ref_node()[1].ref_op_type()), "SequenceLength");
    // No graph inputs (SequenceEmpty takes none and SequenceLength's input
    // is produced in-graph).
    EXPECT_EQ(graph.ref_input().size(), 0u) << tc->name;
    ASSERT_EQ(graph.ref_output().size(), 1u);
    ASSERT_EQ(tc->data_sets().size(), 1u);
    EXPECT_TRUE(tc->data_sets()[0].inputs.empty());
    ASSERT_EQ(tc->data_sets()[0].outputs.size(), 1u);
    const core::runtime::Tensor &out = tc->data_sets()[0].outputs[0];
    EXPECT_EQ(out.data_type, static_cast<int32_t>(core::runtime::DataType::INT64));
    EXPECT_TRUE(out.shape.empty());
    ASSERT_EQ(out.data.size(), sizeof(int64_t));
    EXPECT_EQ(*out.AsInt64(), 0);
  }

  // The default case has no 'dtype' attribute; the int64 case sets it to INT64.
  EXPECT_EQ(default_case->model().ref_graph().ref_node()[0].ref_attribute().size(), 0u);
  ASSERT_EQ(int64_case->model().ref_graph().ref_node()[0].ref_attribute().size(), 1u);
  const AttributeProto &dtype_attr =
      int64_case->model().ref_graph().ref_node()[0].ref_attribute()[0];
  EXPECT_EQ(dtype_attr.ref_name(), "dtype");
  EXPECT_EQ(dtype_attr.ref_type(), AttributeProto::AttributeType::INT);
  EXPECT_EQ(dtype_attr.ref_i(), static_cast<int64_t>(TensorProto::INT64));
}

TEST(BackendTestCase, SplitToSequenceCasesAreRegistered) {
  auto cases = CollectTestCases("SplitToSequence");
  const TestCase *case1 = nullptr;
  const TestCase *case2 = nullptr;
  const TestCase *case_nokd = nullptr;
  for (const auto &c : cases) {
    if (c.name == "test_cc_split_to_sequence_1") {
      case1 = &c;
    } else if (c.name == "test_cc_split_to_sequence_2") {
      case2 = &c;
    } else if (c.name == "test_cc_split_to_sequence_nokeepdims") {
      case_nokd = &c;
    }
  }
  ASSERT_NE(case1, nullptr);
  ASSERT_NE(case2, nullptr);
  ASSERT_NE(case_nokd, nullptr);

  // All cases have a single SplitToSequence node and a sequence-typed output.
  for (const TestCase *tc : {case1, case2, case_nokd}) {
    const GraphProto &graph = tc->model().ref_graph();
    ASSERT_EQ(graph.ref_node().size(), 1u) << tc->name;
    EXPECT_EQ(std::string(graph.ref_node()[0].ref_op_type()), "SplitToSequence") << tc->name;
    ASSERT_EQ(graph.ref_output().size(), 1u);
    EXPECT_TRUE(graph.ref_output()[0].ref_type().has_sequence_type()) << tc->name;
    ASSERT_EQ(tc->data_sets().size(), 1u);
  }

  // Case 1: input [3,6] + scalar split=2 → 3 chunks of [3,2] stacked to [3,3,2].
  EXPECT_EQ(case1->data_sets()[0].inputs.size(), 2u);
  EXPECT_EQ(case1->data_sets()[0].outputs[0].shape, (std::vector<int64_t>{3, 3, 2}));

  // Case 2: input [4,6] + vector split=[2,2] → 2 chunks of [2,6] stacked to [2,2,6].
  EXPECT_EQ(case2->data_sets()[0].inputs.size(), 2u);
  EXPECT_EQ(case2->data_sets()[0].outputs[0].shape, (std::vector<int64_t>{2, 2, 6}));

  // No-keepdims case: input [3,6], keepdims=0 → 6 chunks of [3] stacked to [6,3].
  EXPECT_EQ(case_nokd->data_sets()[0].inputs.size(), 1u);
  EXPECT_EQ(case_nokd->data_sets()[0].outputs[0].shape, (std::vector<int64_t>{6, 3}));
}

} // namespace Test
