TransposeLayer

include/layers/TransposeLayer.hpp

@@ -0,0 +1,33 @@
+#pragma once
+#include <vector>
+
+#include "layers/Layer.hpp"
+#include "layers/Tensor.hpp"
+
+namespace it_lab_ai {
+
+class TransposeLayer : public Layer {
+ public:
+  explicit TransposeLayer(std::vector<int64_t> perm = {})
+      : perm_(std::move(perm)) {}
+
+  void run(const Tensor& input, Tensor& output) override;
+
+#ifdef ENABLE_STATISTIC_WEIGHTS
+  Tensor get_weights() override { return Tensor(); }
+#endif
+
+  static std::string get_name() { return "TransposeLayer"; }
+
+ private:
+  std::vector<int64_t> perm_;
+
+  static void validate_perm(const Shape& input_shape,
+                            const std::vector<int64_t>& perm);
+
+  template <typename T>
+  void transpose_impl(const Tensor& input, Tensor& output,
+                      const std::vector<int64_t>& perm) const;
+};
+
+}  // namespace it_lab_ai

src/layers/TransposeLayer.cpp

@@ -0,0 +1,141 @@
+#include "layers/TransposeLayer.hpp"
+
+#include <algorithm>
+#include <numeric>
+
+namespace it_lab_ai {
+
+void TransposeLayer::run(const Tensor& input, Tensor& output) {
+  const auto& shape = input.get_shape();
+
+  std::vector<int64_t> perm = perm_;
+  if (perm.empty()) {
+    perm.resize(shape.dims());
+    std::iota(perm.begin(), perm.end(), 0);
+  }
+
+  validate_perm(shape, perm);
+
+  switch (input.get_type()) {
+    case Type::kFloat:
+      transpose_impl<float>(input, output, perm);
+      break;
+    case Type::kInt:
+      transpose_impl<int>(input, output, perm);
+      break;
+    default:
+      throw std::runtime_error("Unsupported tensor data type");
+  }
+}
+
+template <typename T>
+void TransposeLayer::transpose_impl(const Tensor& input, Tensor& output,
+                                    const std::vector<int64_t>& perm) const {
+  const auto& shape = input.get_shape();
+  const auto* input_data = input.as<T>();
+
+  if (!input_data || input_data->empty()) {
+    throw std::runtime_error("Input tensor is empty or invalid");
+  }
+
+  std::vector<size_t> new_dims;
+  new_dims.reserve(shape.dims());
+  for (const auto& axis : perm) {
+    new_dims.push_back(shape[static_cast<size_t>(axis)]);
+  }
+  Shape new_shape(new_dims);
+
+  std::vector<size_t> input_strides(shape.dims());
+  size_t stride = 1;
+  for (size_t dim = shape.dims(); dim-- > 0;) {
+    input_strides[dim] = stride;
+    stride *= shape[dim];
+  }
+
+  std::vector<size_t> output_strides(new_shape.dims());
+  stride = 1;
+  for (size_t dim = new_shape.dims(); dim-- > 0;) {
+    output_strides[dim] = stride;
+    stride *= new_shape[dim];
+  }
+
+  std::vector<T> output_values(input_data->size());
+
+  if (shape.dims() == 2) {
+    const size_t rows = shape[0];
+    const size_t cols = shape[1];
+
+    if (perm[0] == 1 && perm[1] == 0) {
+      for (size_t i = 0; i < rows; ++i) {
+        for (size_t j = 0; j < cols; ++j) {
+          const size_t new_index = j * rows + i;
+          if (new_index >= output_values.size()) {
+            throw std::runtime_error(
+                "Index out of bounds during 2D transposition");
+          }
+          output_values[new_index] = (*input_data)[i * cols + j];
+        }
+      }
+    } else {
+      for (size_t i = 0; i < input_data->size(); ++i) {
+        size_t old_index = i;
+        size_t new_index = 0;
+
+        for (size_t dim = 0; dim < perm.size(); ++dim) {
+          const auto axis = static_cast<size_t>(perm[dim]);
+          const size_t coord = (old_index / input_strides[axis]) % shape[axis];
+          new_index += coord * output_strides[dim];
+        }
+
+        if (new_index >= output_values.size()) {
+          throw std::runtime_error(
+              "Index out of bounds during 2D transposition");
+        }
+        output_values[new_index] = (*input_data)[i];
+      }
+    }
+  } else {
+    for (size_t i = 0; i < input_data->size(); ++i) {
+      size_t old_index = i;
+      size_t new_index = 0;
+
+      for (size_t dim = 0; dim < perm.size(); ++dim) {
+        const auto axis = static_cast<size_t>(perm[dim]);
+        const size_t coord = (old_index / input_strides[axis]) % shape[axis];
+        new_index += coord * output_strides[dim];
+      }
+
+      if (new_index >= output_values.size()) {
+        throw std::runtime_error("Index out of bounds during transposition");
+      }
+      output_values[new_index] = (*input_data)[i];
+    }
+  }
+
+  output = make_tensor(output_values, new_shape);
+}
+
+void TransposeLayer::validate_perm(const Shape& input_shape,
+                                   const std::vector<int64_t>& perm) {
+  if (perm.size() != input_shape.dims()) {
+    throw std::invalid_argument("Permutation size must match input dimensions");
+  }
+
+  std::vector<bool> used_axes(input_shape.dims(), false);
+  for (const auto& axis : perm) {
+    if (axis < 0 || static_cast<size_t>(axis) >= input_shape.dims()) {
+      throw std::invalid_argument("Invalid axis in permutation");
+    }
+    if (used_axes[static_cast<size_t>(axis)]) {
+      throw std::invalid_argument("Duplicate axis in permutation");
+    }
+    used_axes[static_cast<size_t>(axis)] = true;
+  }
+}
+
+template void TransposeLayer::transpose_impl<float>(
+    const Tensor&, Tensor&, const std::vector<int64_t>&) const;
+template void TransposeLayer::transpose_impl<int>(
+    const Tensor&, Tensor&, const std::vector<int64_t>&) const;
+
+}  // namespace it_lab_ai

test/single_layer/test_transposelayer.cpp

@@ -0,0 +1,199 @@
+#include <cstdint>
+#include <iostream>
+#include <vector>
+
+#include "gtest/gtest.h"
+#include "layers/Tensor.hpp"
+#include "layers/TransposeLayer.hpp"
+
+using namespace it_lab_ai;
+
+TEST(TransposeLayerTest, EmptyTensor) {
+  Tensor input = make_tensor<float>({}, {0});
+  TransposeLayer layer;
+  Tensor output;
+
+  EXPECT_THROW(layer.run(input, output), std::runtime_error);
+}
+
+TEST(TransposeLayerTest, IdentityTranspose) {
+  Tensor input = make_tensor<float>({1, 2, 3, 4}, {2, 2});
+  TransposeLayer layer({0, 1});
+  Tensor output;
+
+  layer.run(input, output);
+
+  ASSERT_EQ(output.get_shape(), Shape({2, 2}));
+  EXPECT_FLOAT_EQ(output.get<float>({0, 0}), 1.0f);
+  EXPECT_FLOAT_EQ(output.get<float>({0, 1}), 2.0f);
+  EXPECT_FLOAT_EQ(output.get<float>({1, 0}), 3.0f);
+  EXPECT_FLOAT_EQ(output.get<float>({1, 1}), 4.0f);
+}
+
+TEST(TransposeLayerTest, VectorTranspose) {
+  Tensor input = make_tensor<float>({1, 2, 3, 4}, {4});
+  TransposeLayer layer({0});
+  Tensor output;
+
+  layer.run(input, output);
+
+  ASSERT_EQ(output.get_shape(), Shape({4}));
+  EXPECT_FLOAT_EQ(output.get<float>({0}), 1.0f);
+  EXPECT_FLOAT_EQ(output.get<float>({1}), 2.0f);
+  EXPECT_FLOAT_EQ(output.get<float>({2}), 3.0f);
+  EXPECT_FLOAT_EQ(output.get<float>({3}), 4.0f);
+}
+
+TEST(TransposeLayerTest, InvalidPermutationSize) {
+  Tensor input = make_tensor<float>({1, 2, 3, 4}, {2, 2});
+  TransposeLayer layer({0});
+  Tensor output;
+
+  EXPECT_THROW(layer.run(input, output), std::invalid_argument);
+}
+
+TEST(TransposeLayerTest, DuplicateAxes) {
+  Tensor input = make_tensor<float>({1, 2, 3, 4}, {2, 2});
+  TransposeLayer layer({0, 0});
+  Tensor output;
+
+  EXPECT_THROW(layer.run(input, output), std::invalid_argument);
+}
+
+TEST(TransposeLayerTest, NegativeAxis) {
+  Tensor input = make_tensor<float>({1, 2, 3, 4}, {2, 2});
+  TransposeLayer layer({0, -1});
+  Tensor output;
+
+  EXPECT_THROW(layer.run(input, output), std::invalid_argument);
+}
+
+TEST(TransposeLayerTest, LargeAxis) {
+  Tensor input = make_tensor<float>({1, 2, 3, 4}, {2, 2});
+  TransposeLayer layer({0, 2});
+  Tensor output;
+
+  EXPECT_THROW(layer.run(input, output), std::invalid_argument);
+}
+
+TEST(TransposeLayerTest, 4DTensorTranspose) {
+  std::vector<float> data(16);
+  std::iota(data.begin(), data.end(), 1.0f);
+  Tensor input = make_tensor<float>(data, {2, 2, 2, 2});
+  TransposeLayer layer({3, 1, 0, 2});
+  Tensor output;
+
+  layer.run(input, output);
+
+  ASSERT_EQ(output.get_shape(), Shape({2, 2, 2, 2}));
+  EXPECT_FLOAT_EQ(output.get<float>({0, 0, 0, 0}), 1.0f);
+  EXPECT_FLOAT_EQ(output.get<float>({1, 0, 0, 0}), 2.0f);
+  EXPECT_FLOAT_EQ(output.get<float>({0, 1, 0, 0}), 5.0f);
+  EXPECT_FLOAT_EQ(output.get<float>({1, 1, 0, 0}), 6.0f);
+  EXPECT_FLOAT_EQ(output.get<float>({0, 0, 0, 1}), 3.0f);
+  EXPECT_FLOAT_EQ(output.get<float>({1, 0, 0, 1}), 4.0f);
+  EXPECT_FLOAT_EQ(output.get<float>({0, 1, 0, 1}), 7.0f);
+  EXPECT_FLOAT_EQ(output.get<float>({1, 1, 0, 1}), 8.0f);
+  EXPECT_FLOAT_EQ(output.get<float>({0, 0, 1, 0}), 9.0f);
+  EXPECT_FLOAT_EQ(output.get<float>({1, 0, 1, 0}), 10.0f);
+  EXPECT_FLOAT_EQ(output.get<float>({0, 1, 1, 0}), 13.0f);
+  EXPECT_FLOAT_EQ(output.get<float>({1, 1, 1, 0}), 14.0f);
+  EXPECT_FLOAT_EQ(output.get<float>({0, 0, 1, 1}), 11.0f);
+  EXPECT_FLOAT_EQ(output.get<float>({1, 0, 1, 1}), 12.0f);
+  EXPECT_FLOAT_EQ(output.get<float>({0, 1, 1, 1}), 15.0f);
+  EXPECT_FLOAT_EQ(output.get<float>({1, 1, 1, 1}), 16.0f);
+}
+
+TEST(TransposeLayerTest, MatrixTranspose) {
+  Tensor input = make_tensor<float>({1, 2, 3, 4}, {2, 2});
+  TransposeLayer layer({1, 0});
+  Tensor output;
+
+  layer.run(input, output);
+
+  ASSERT_EQ(output.get_shape(), Shape({2, 2}));
+  EXPECT_FLOAT_EQ(output.get<float>({0, 0}), 1.0f);
+  EXPECT_FLOAT_EQ(output.get<float>({0, 1}), 3.0f);
+  EXPECT_FLOAT_EQ(output.get<float>({1, 0}), 2.0f);
+  EXPECT_FLOAT_EQ(output.get<float>({1, 1}), 4.0f);
+}
+
+TEST(TransposeLayerTest, 3DTensor) {
+  std::vector<float> data(24);
+  std::iota(data.begin(), data.end(), 1.0f);
+  Tensor input = make_tensor<float>(data, {2, 3, 4});
+  TransposeLayer layer({2, 0, 1});
+  Tensor output;
+
+  layer.run(input, output);
+
+  ASSERT_EQ(output.get_shape(), Shape({4, 2, 3}));
+  EXPECT_FLOAT_EQ(output.get<float>({0, 0, 0}), 1.0f);
+  EXPECT_FLOAT_EQ(output.get<float>({3, 1, 2}), 24.0f);
+}
+
+TEST(TransposeLayerTest, IntTensor) {
+  Tensor input = make_tensor<int>({1, 2, 3, 4}, {2, 2});
+  TransposeLayer layer({1, 0});
+  Tensor output;
+  layer.run(input, output);
+  EXPECT_EQ(output.get<int>({0, 0}), 1);
+  EXPECT_EQ(output.get<int>({1, 0}), 2);
+  EXPECT_EQ(output.get<int>({0, 1}), 3);
+  EXPECT_EQ(output.get<int>({1, 1}), 4);
+}
+
+TEST(TransposeLayerTest, 1DDefaultPermutationIsNoOp) {
+  Tensor input = make_tensor<float>({1.0f, 2.0f, 3.0f, 4.0f}, {4});
+
+  TransposeLayer layer;
+  Tensor output;
+
+  EXPECT_NO_THROW(layer.run(input, output));
+
+  EXPECT_EQ(output.get_shape(), Shape({4}));
+  EXPECT_FLOAT_EQ(output.get<float>({0}), 1.0f);
+  EXPECT_FLOAT_EQ(output.get<float>({1}), 2.0f);
+  EXPECT_FLOAT_EQ(output.get<float>({2}), 3.0f);
+  EXPECT_FLOAT_EQ(output.get<float>({3}), 4.0f);
+}
+
+TEST(TransposeLayerTest, MultipleRunsWithDifferentRanks) {
+  TransposeLayer layer({1, 0});
+
+  {
+    Tensor input = make_tensor<float>({1, 2, 3, 4}, {2, 2});
+    Tensor output;
+    layer.run(input, output);
+    EXPECT_EQ(output.get_shape(), Shape({2, 2}));
+    EXPECT_FLOAT_EQ(output.get<float>({0, 0}), 1.0f);
+    EXPECT_FLOAT_EQ(output.get<float>({1, 0}), 2.0f);
+  }
+
+  {
+    Tensor input = make_tensor<float>({1, 2, 3, 4, 5, 6}, {3, 2});
+    Tensor output;
+    layer.run(input, output);
+    EXPECT_EQ(output.get_shape(), Shape({2, 3}));
+  }
+
+  {
+    Tensor input = make_tensor<float>({1, 2, 3, 4, 5, 6, 7, 8}, {2, 2, 2});
+    Tensor output;
+    EXPECT_THROW(layer.run(input, output), std::invalid_argument);
+  }
+}
+
+TEST(TransposeLayerTest, ExplicitPermutationWithDifferentRanks) {
+  TransposeLayer layer({1, 0});
+
+  Tensor input2D = make_tensor<float>({1, 2, 3, 4}, {2, 2});
+  Tensor output2D;
+  layer.run(input2D, output2D);
+  EXPECT_EQ(output2D.get_shape(), Shape({2, 2}));
+  EXPECT_FLOAT_EQ(output2D.get<float>({1, 0}), 2.0f);
+
+  Tensor input3D = make_tensor<float>({1, 2, 3, 4, 5, 6, 7, 8}, {2, 2, 2});
+  Tensor output3D;
+  EXPECT_THROW(layer.run(input3D, output3D), std::invalid_argument);
+}