SplitLayer

include/layers/SplitLayer.hpp

@@ -0,0 +1,38 @@
+#pragma once
+#include <optional>
+#include <stdexcept>
+#include <vector>
+
+#include "layers/Layer.hpp"
+#include "layers/Tensor.hpp"
+
+namespace it_lab_ai {
+
+class SplitLayer : public Layer {
+ public:
+  SplitLayer(int axis, std::vector<int> splits)
+      : axis_(axis), splits_(std::move(splits)) {}
+
+  SplitLayer(int axis, int num_outputs)
+      : axis_(axis), num_outputs_(num_outputs) {}
+  void run(const Tensor& input, Tensor& output) override;
+  void run(const Tensor& input, std::vector<Tensor>& outputs);
+
+  static std::string get_name() { return "SplitLayer"; }
+
+#ifdef ENABLE_STATISTIC_WEIGHTS
+  Tensor get_weights() override { return Tensor(); }
+#endif
+
+ private:
+  int axis_;
+  std::optional<std::vector<int>> splits_;
+  std::optional<int> num_outputs_;
+
+  void validate(const Tensor& input) const;
+  int get_normalized_axis(int rank) const;
+  template <typename T>
+  void split_impl(const Tensor& input, std::vector<Tensor>& outputs) const;
+};
+
+}  // namespace it_lab_ai

src/layers/SplitLayer.cpp

@@ -0,0 +1,131 @@
+#include "layers/SplitLayer.hpp"
+
+#include <algorithm>
+#include <cstring>
+
+namespace it_lab_ai {
+
+void SplitLayer::run(const Tensor& input, Tensor& output) { output = input; }
+
+void SplitLayer::run(const Tensor& input, std::vector<Tensor>& outputs) {
+  validate(input);
+
+  switch (input.get_type()) {
+    case Type::kFloat:
+      split_impl<float>(input, outputs);
+      break;
+    case Type::kInt:
+      split_impl<int>(input, outputs);
+      break;
+    default:
+      throw std::runtime_error("Unsupported tensor data type");
+  }
+}
+
+template <typename T>
+void SplitLayer::split_impl(const Tensor& input,
+                            std::vector<Tensor>& outputs) const {
+  const auto& input_data = *input.as<T>();
+  const Shape& shape = input.get_shape();
+  const int axis = get_normalized_axis(static_cast<int>(shape.dims()));
+
+  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);
+    }
+  }
+
+  size_t outer_size = 1;
+  for (int i = 0; i < axis; ++i) {
+    outer_size *= shape[i];
+  }
+
+  size_t inner_size = 1;
+  for (size_t i = axis + 1; i < shape.dims(); ++i) {
+    inner_size *= shape[i];
+  }
+
+  const size_t input_axis_stride = shape[axis] * inner_size;
+
+  outputs.clear();
+  outputs.reserve(part_sizes.size());
+
+  size_t input_offset = 0;
+  for (const auto part_size : part_sizes) {
+    const auto output_axis_size = static_cast<size_t>(part_size);
+
+    std::vector<size_t> output_shape_vec(shape.dims());
+    for (size_t i = 0; i < shape.dims(); ++i) {
+      output_shape_vec[i] =
+          (static_cast<int>(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>();
+
+    const size_t output_part_size = output_axis_size * inner_size;
+
+    for (size_t outer = 0; outer < outer_size; ++outer) {
+      const T* input_start =
+          &input_data[outer * input_axis_stride + input_offset * inner_size];
+      T* output_start = &output_data[outer * output_part_size];
+
+      std::copy_n(input_start, output_part_size, output_start);
+    }
+
+    input_offset += output_axis_size;
+  }
+}
+
+void SplitLayer::validate(const Tensor& input) const {
+  if (input.get_shape().dims() == 0) {
+    throw std::runtime_error("Cannot split scalar tensor");
+  }
+
+  const int axis =
+      get_normalized_axis(static_cast<int>(input.get_shape().dims()));
+  const int axis_size = static_cast<int>(input.get_shape()[axis]);
+
+  if (splits_) {
+    int sum = 0;
+    for (int s : *splits_) {
+      if (s <= 0) throw std::runtime_error("Split size must be positive");
+      sum += s;
+    }
+    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");
+    }
+    if (*num_outputs_ > axis_size) {
+      throw std::runtime_error("num_outputs (" + std::to_string(*num_outputs_) +
+                               ") cannot be greater than axis size (" +
+                               std::to_string(axis_size) + ")");
+    }
+  }
+}
+
+int SplitLayer::get_normalized_axis(int rank) const {
+  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&,
+                                            std::vector<Tensor>&) const;
+template void SplitLayer::split_impl<int>(const Tensor&,
+                                          std::vector<Tensor>&) const;
+
+}  // namespace it_lab_ai

test/single_layer/test_splitlayer.cpp

@@ -0,0 +1,215 @@
+#include <vector>
+
+#include "gtest/gtest.h"
+#include "layers/SplitLayer.hpp"
+#include "layers/Tensor.hpp"
+
+using namespace it_lab_ai;
+
+TEST(SplitLayerTests, SplitEqualParts1D) {
+  Tensor input = make_tensor<float>({1, 2, 3, 4, 5, 6}, {6});
+  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({2}));
+  EXPECT_FLOAT_EQ(outputs[0].get<float>({0}), 1.0f);
+  EXPECT_FLOAT_EQ(outputs[1].get<float>({0}), 3.0f);
+  EXPECT_FLOAT_EQ(outputs[2].get<float>({0}), 5.0f);
+}
+
+TEST(SplitLayerTests, SplitVariableParts1D) {
+  Tensor input = make_tensor<float>({1, 2, 3, 4, 5, 6}, {6});
+  SplitLayer splitter(0, {2, 4});
+
+  std::vector<Tensor> outputs;
+  splitter.run(input, outputs);
+
+  ASSERT_EQ(outputs.size(), 2);
+  EXPECT_EQ(outputs[0].get_shape(), Shape({2}));
+  EXPECT_EQ(outputs[1].get_shape(), Shape({4}));
+  EXPECT_FLOAT_EQ(outputs[0].get<float>({1}), 2.0f);
+  EXPECT_FLOAT_EQ(outputs[1].get<float>({3}), 6.0f);
+}
+
+TEST(SplitLayerTests, Split2DAlongAxis0) {
+  Tensor input = make_tensor<float>({1, 2, 3, 4, 5, 6}, {2, 3});
+  SplitLayer splitter(0, {1, 1});
+
+  std::vector<Tensor> outputs;
+  splitter.run(input, outputs);
+
+  ASSERT_EQ(outputs.size(), 2);
+  EXPECT_EQ(outputs[0].get_shape(), Shape({1, 3}));
+  EXPECT_EQ(outputs[1].get_shape(), Shape({1, 3}));
+  EXPECT_FLOAT_EQ(outputs[0].get<float>({0, 2}), 3.0f);
+  EXPECT_FLOAT_EQ(outputs[1].get<float>({0, 0}), 4.0f);
+}
+
+TEST(SplitLayerTests, Split2DAlongAxis1) {
+  Tensor input = make_tensor<float>({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_FLOAT_EQ(outputs[0].get<float>({1, 0}), 4.0f);
+  EXPECT_FLOAT_EQ(outputs[1].get<float>({0, 1}), 3.0f);
+}
+
+TEST(SplitLayerTests, Split3DEqualParts) {
+  std::vector<float> data(2 * 3 * 4);
+  std::iota(data.begin(), data.end(), 0.0f);
+  Tensor input = make_tensor<float>(data, {2, 3, 4});
+
+  SplitLayer splitter(1, 3);
+
+  std::vector<Tensor> outputs;
+  splitter.run(input, outputs);
+
+  ASSERT_EQ(outputs.size(), 3);
+  EXPECT_EQ(outputs[0].get_shape(), Shape({2, 1, 4}));
+  EXPECT_EQ(outputs[1].get<float>({1, 0, 3}), 19.0f);
+}
+
+TEST(SplitLayerTests, Split4DVariableParts) {
+  std::vector<float> data(1 * 3 * 2 * 4);
+  std::iota(data.begin(), data.end(), 0.0f);
+  Tensor input = make_tensor<float>(data, {1, 3, 2, 4});
+
+  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, 3, 1, 4}));
+  EXPECT_EQ(outputs[1].get<float>({0, 2, 0, 3}), 23.0f);
+}
+
+TEST(SplitLayerTests, SplitNegativeAxis) {
+  Tensor input = make_tensor<float>({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}));
+}
+
+TEST(SplitLayerTests, InvalidSplitSizes) {
+  Tensor input = make_tensor<float>({1, 2, 3, 4}, {4});
+
+  SplitLayer splitter(0, {1, 2});
+
+  std::vector<Tensor> outputs;
+  EXPECT_THROW(splitter.run(input, outputs), std::runtime_error);
+}
+
+TEST(SplitLayerTests, EmptyInputTensor) {
+  Tensor input = make_tensor<float>({}, {0});
+
+  SplitLayer splitter(0, {});
+
+  std::vector<Tensor> outputs;
+  EXPECT_THROW(splitter.run(input, outputs), std::runtime_error);
+}
+
+TEST(SplitLayerTests, Split192IntoTwo96) {
+  std::vector<float> input_data(1 * 192 * 56 * 56);
+  std::iota(input_data.begin(), input_data.end(), 0.0f);
+  Tensor input = make_tensor<float>(input_data, {1, 192, 56, 56});
+
+  SplitLayer splitter(1, {96, 96});
+  std::vector<Tensor> outputs;
+  splitter.run(input, outputs);
+
+  ASSERT_EQ(outputs.size(), 2);
+  EXPECT_EQ(outputs[0].get_shape(), Shape({1, 96, 56, 56}));
+  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);
+}