SplitLayer

include/layers/SplitLayer.hpp

@@ -0,0 +1,36 @@
+#pragma once
+#include <vector>
+
+#include "layers/Layer.hpp"
+#include "layers/Tensor.hpp"
+
+namespace it_lab_ai {
+
+class SplitLayer : public Layer {
+ public:
+  SplitLayer(int axis, const std::vector<int>& splits)
+      : axis_(axis), splits_(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::vector<int> splits_;
+  int num_outputs_ = 0;
+
+  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,119 @@
+#include "layers/SplitLayer.hpp"
+
+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);
+  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:
+      split_impl<float>(input, outputs);
+      break;
+    case Type::kInt:
+      split_impl<int>(input, outputs);
+      break;
+    default:
+      throw std::runtime_error("Unsupported tensor 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()));
+  const auto& part_sizes =
+      splits_.empty()
+          ? std::vector<int>(num_outputs_,
+                             static_cast<int>(shape[axis]) / num_outputs_)
+          : splits_;
+
+  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];
+  }
+
+  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 outer = 0; outer < outer_size; ++outer) {
+      for (size_t a = 0; a < output_axis_size; ++a) {
+        for (size_t inner = 0; inner < inner_size; ++inner) {
+          size_t input_pos = outer * shape[axis] * inner_size +
+                             (input_offset + a) * inner_size + inner;
+          size_t output_pos =
+              outer * output_axis_size * inner_size + a * inner_size + inner;
+          output_data[output_pos] = input_data[input_pos];
+        }
+      }
+    }
+    input_offset += output_axis_size;
+  }
+}
+
+void SplitLayer::validate(const Tensor& input) const {
+  if (input.get_shape().dims() == 0) {
+    throw std::runtime_error("SplitLayer: 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];
+
+  if (!splits_.empty()) {
+    int sum = 0;
+    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)) {
+      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");
+  }
+}
+
+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_;
+}
+
+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,141 @@
+#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);
+}