change to optional, add tests, fix check with large negative axis

Author: Semyon1104

include/layers/SplitLayer.hpp

@@ -1,4 +1,6 @@
 #pragma once
+#include <optional>
+#include <stdexcept>
 #include <vector>
 
 #include "layers/Layer.hpp"
@@ -8,9 +10,10 @@ namespace it_lab_ai {
 
 class SplitLayer : public Layer {
  public:
-  SplitLayer(int axis, const std::vector<int>& splits)
-      : axis_(axis), splits_(splits) {}
+  SplitLayer(int axis, std::vector<int> splits)
+      : axis_(axis), splits_(std::move(splits)) {}
 
+  // Ðåæèì 2: êîëè÷åñòâî âûõîäíûõ òåíçîðîâ
   SplitLayer(int axis, int num_outputs)
       : axis_(axis), num_outputs_(num_outputs) {}
   void run(const Tensor& input, Tensor& output) override;
@@ -24,8 +27,8 @@ class SplitLayer : public Layer {
 
  private:
   int axis_;
-  std::vector<int> splits_;
-  int num_outputs_ = 0;
+  std::optional<std::vector<int>> splits_;
+  std::optional<int> num_outputs_;
 
   void validate(const Tensor& input) const;
   int get_normalized_axis(int rank) const;

src/layers/SplitLayer.cpp

@@ -6,27 +6,6 @@ void SplitLayer::run(const Tensor& input, Tensor& output) { output = input; }
 
 void SplitLayer::run(const Tensor& input, std::vector<Tensor>& outputs) {
   validate(input);
-  const auto& shape = input.get_shape();
-  const int axis = get_normalized_axis(static_cast<int>(shape.dims()));
-
-  std::vector<int> part_sizes;
-  if (!splits_.empty()) {
-    part_sizes = splits_;
-  } else {
-    const int base_size = static_cast<int>(shape[axis]) / num_outputs_;
-    const int remainder = static_cast<int>(shape[axis]) % num_outputs_;
-    part_sizes.assign(num_outputs_, base_size);
-    if (remainder > 0) {
-      part_sizes.back() += remainder;
-    }
-  }
-
-  outputs.clear();
-  for (int size : part_sizes) {
-    Shape out_shape = shape;
-    out_shape[axis] = static_cast<size_t>(size);
-    outputs.emplace_back(out_shape, input.get_type());
-  }
 
   switch (input.get_type()) {
     case Type::kFloat:
@@ -36,7 +15,7 @@ void SplitLayer::run(const Tensor& input, std::vector<Tensor>& outputs) {
       split_impl<int>(input, outputs);
       break;
     default:
-      throw std::runtime_error("Unsupported tensor type");
+      throw std::runtime_error("Unsupported tensor data type");
   }
 }
 
@@ -46,11 +25,23 @@ void SplitLayer::split_impl(const Tensor& input,
   const auto& input_data = *input.as<T>();
   const Shape& shape = input.get_shape();
   const int axis = get_normalized_axis(static_cast<int>(shape.dims()));
-  const auto& part_sizes =
-      splits_.empty()
-          ? std::vector<int>(num_outputs_,
-                             static_cast<int>(shape[axis]) / num_outputs_)
-          : splits_;
+
+  std::vector<int> part_sizes;
+  if (splits_) {
+    part_sizes = *splits_;
+  } else {
+    const int total_size = static_cast<int>(shape[axis]);
+    const int base_size = total_size / *num_outputs_;
+    const int remainder = total_size % *num_outputs_;
+
+    part_sizes.reserve(*num_outputs_);
+    for (int i = 0; i < *num_outputs_; ++i) {
+      part_sizes.push_back(i < remainder ? base_size + 1 : base_size);
+    }
+  }
+
+  outputs.clear();
+  outputs.reserve(part_sizes.size());
 
   size_t outer_size = 1;
   for (int i = 0; i < axis; ++i) {
@@ -63,9 +54,17 @@ void SplitLayer::split_impl(const Tensor& input,
   }
 
   size_t input_offset = 0;
-  for (auto& output : outputs) {
-    auto& output_data = *output.as<T>();
-    const size_t output_axis_size = output.get_shape()[axis];
+  for (size_t part = 0; part < part_sizes.size(); ++part) {
+    const size_t output_axis_size = part_sizes[part];
+
+    std::vector<size_t> output_shape_vec(shape.dims());
+    for (size_t i = 0; i < shape.dims(); ++i) {
+      output_shape_vec[i] = (i == axis) ? output_axis_size : shape[i];
+    }
+    Shape output_shape(output_shape_vec);
+
+    outputs.emplace_back(output_shape, input.get_type());
+    auto& output_data = *outputs.back().as<T>();
 
     for (size_t outer = 0; outer < outer_size; ++outer) {
       for (size_t a = 0; a < output_axis_size; ++a) {
@@ -84,31 +83,37 @@ void SplitLayer::split_impl(const Tensor& input,
 
 void SplitLayer::validate(const Tensor& input) const {
   if (input.get_shape().dims() == 0) {
-    throw std::runtime_error("SplitLayer: Cannot split scalar tensor");
+    throw std::runtime_error("Cannot split scalar tensor");
   }
 
   const int axis =
       get_normalized_axis(static_cast<int>(input.get_shape().dims()));
-  const size_t axis_size = input.get_shape()[axis];
+  const int axis_size = static_cast<int>(input.get_shape()[axis]);
 
-  if (!splits_.empty()) {
+  if (splits_) {
     int sum = 0;
-    for (int s : splits_) {
+    for (int s : *splits_) {
       if (s <= 0) throw std::runtime_error("Split size must be positive");
       sum += s;
     }
-    if (sum != static_cast<int>(axis_size)) {
+    if (sum != axis_size) {
       throw std::runtime_error("Sum of splits must match axis size");
     }
-  } else if (num_outputs_ <= 0) {
-    throw std::runtime_error("num_outputs must be positive");
+  } else {
+    if (*num_outputs_ <= 0) {
+      throw std::runtime_error("num_outputs must be positive");
+    }
+    if (*num_outputs_ > axis_size) {
+      throw std::runtime_error("num_outputs cannot be greater than axis size");
+    }
   }
 }
 
 int SplitLayer::get_normalized_axis(int rank) const {
-  if (axis_ < 0) return axis_ + rank;
-  if (axis_ >= rank) throw std::runtime_error("Axis out of bounds");
-  return axis_;
+  if (axis_ < -rank || axis_ >= rank) {
+    throw std::runtime_error("Axis out of bounds");
+  }
+  return (axis_ < 0) ? axis_ + rank : axis_;
 }
 
 template void SplitLayer::split_impl<float>(const Tensor&,

test/single_layer/test_splitlayer.cpp

@@ -138,4 +138,79 @@ TEST(SplitLayerTests, Split192IntoTwo96) {
   EXPECT_EQ(outputs[1].get_shape(), Shape({1, 96, 56, 56}));
   EXPECT_FLOAT_EQ(outputs[0].get<float>({0, 0, 0, 0}), 0.0f);
   EXPECT_FLOAT_EQ(outputs[1].get<float>({0, 0, 0, 0}), 96 * 56 * 56);
+}
+
+TEST(SplitLayerTests, UnevenSplitWithRemainder) {
+  Tensor input = make_tensor<float>({1, 2, 3, 4, 5}, {5});
+  SplitLayer splitter(0, 3);
+
+  std::vector<Tensor> outputs;
+  splitter.run(input, outputs);
+
+  ASSERT_EQ(outputs.size(), 3);
+  EXPECT_EQ(outputs[0].get_shape(), Shape({2}));
+  EXPECT_EQ(outputs[1].get_shape(), Shape({2}));
+  EXPECT_EQ(outputs[2].get_shape(), Shape({1}));
+  EXPECT_FLOAT_EQ(outputs[0].get<float>({1}), 2.0f);
+  EXPECT_FLOAT_EQ(outputs[1].get<float>({1}), 4.0f);
+  EXPECT_FLOAT_EQ(outputs[2].get<float>({0}),
+                  5.0f);
+}
+
+TEST(SplitLayerTests, NumOutputsGreaterThanAxisSize) {
+  Tensor input = make_tensor<float>({1, 2, 3}, {3});
+  SplitLayer splitter(0, 5);
+
+  std::vector<Tensor> outputs;
+  EXPECT_THROW(splitter.run(input, outputs), std::runtime_error);
+}
+
+TEST(SplitLayerTests, IntegerDataType) {
+  Tensor input = make_tensor<int>({1, 2, 3, 4, 5, 6}, {2, 3});
+  SplitLayer splitter(1, {1, 2});
+
+  std::vector<Tensor> outputs;
+  splitter.run(input, outputs);
+
+  ASSERT_EQ(outputs.size(), 2);
+  EXPECT_EQ(outputs[0].get_shape(), Shape({2, 1}));
+  EXPECT_EQ(outputs[1].get_shape(), Shape({2, 2}));
+  EXPECT_EQ(outputs[0].get<int>({1, 0}), 4);
+  EXPECT_EQ(outputs[1].get<int>({0, 1}), 3);
+}
+
+TEST(SplitLayerTests, NegativeAxis2D) {
+  Tensor input = make_tensor<float>({1, 2, 3, 4}, {2, 2});
+  SplitLayer splitter(-2, {1, 1});
+
+  std::vector<Tensor> outputs;
+  splitter.run(input, outputs);
+
+  ASSERT_EQ(outputs.size(), 2);
+  EXPECT_EQ(outputs[0].get_shape(), Shape({1, 2}));
+  EXPECT_EQ(outputs[1].get_shape(), Shape({1, 2}));
+}
+
+TEST(SplitLayerTests, NegativeAxis3D) {
+  std::vector<float> data(2 * 3 * 4);
+  std::iota(data.begin(), data.end(), 1.0f);
+  Tensor input = make_tensor<float>(data, {2, 3, 4});
+
+  SplitLayer splitter(-1, {1, 3});
+
+  std::vector<Tensor> outputs;
+  splitter.run(input, outputs);
+
+  ASSERT_EQ(outputs.size(), 2);
+  EXPECT_EQ(outputs[0].get_shape(), Shape({2, 3, 1}));
+  EXPECT_EQ(outputs[1].get_shape(), Shape({2, 3, 3}));
+  EXPECT_FLOAT_EQ(outputs[0].get<float>({1, 2, 0}), 21.0f);
+}
+
+TEST(SplitLayerTests, LargeAxisValue) {
+  Tensor input = make_tensor<float>({1, 2, 3, 4}, {2, 2});
+
+  SplitLayer splitter(10, {1, 1});
+  std::vector<Tensor> outputs;
+  EXPECT_THROW(splitter.run(input, outputs), std::runtime_error);
 }