Skip to content

Commit acba1bb

Browse files
Semyon1104Semyon1104
authored
TransposeLayer (#193)
1 parent 4e5a6c6 commit acba1bb

3 files changed

Lines changed: 373 additions & 0 deletions

File tree

‎include/layers/TransposeLayer.hpp‎

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,33 @@
1+
#pragma once
2+
#include <vector>
3+
4+
#include "layers/Layer.hpp"
5+
#include "layers/Tensor.hpp"
6+
7+
namespace it_lab_ai {
8+
9+
class TransposeLayer : public Layer {
10+
public:
11+
explicit TransposeLayer(std::vector<int64_t> perm = {})
12+
: perm_(std::move(perm)) {}
13+
14+
void run(const Tensor& input, Tensor& output) override;
15+
16+
#ifdef ENABLE_STATISTIC_WEIGHTS
17+
Tensor get_weights() override { return Tensor(); }
18+
#endif
19+
20+
static std::string get_name() { return "TransposeLayer"; }
21+
22+
private:
23+
std::vector<int64_t> perm_;
24+
25+
static void validate_perm(const Shape& input_shape,
26+
const std::vector<int64_t>& perm);
27+
28+
template <typename T>
29+
void transpose_impl(const Tensor& input, Tensor& output,
30+
const std::vector<int64_t>& perm) const;
31+
};
32+
33+
} // namespace it_lab_ai

‎src/layers/TransposeLayer.cpp‎

Lines changed: 141 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,141 @@
1+
#include "layers/TransposeLayer.hpp"
2+
3+
#include <algorithm>
4+
#include <numeric>
5+
6+
namespace it_lab_ai {
7+
8+
void TransposeLayer::run(const Tensor& input, Tensor& output) {
9+
const auto& shape = input.get_shape();
10+
11+
std::vector<int64_t> perm = perm_;
12+
if (perm.empty()) {
13+
perm.resize(shape.dims());
14+
std::iota(perm.begin(), perm.end(), 0);
15+
}
16+
17+
validate_perm(shape, perm);
18+
19+
switch (input.get_type()) {
20+
case Type::kFloat:
21+
transpose_impl<float>(input, output, perm);
22+
break;
23+
case Type::kInt:
24+
transpose_impl<int>(input, output, perm);
25+
break;
26+
default:
27+
throw std::runtime_error("Unsupported tensor data type");
28+
}
29+
}
30+
31+
template <typename T>
32+
void TransposeLayer::transpose_impl(const Tensor& input, Tensor& output,
33+
const std::vector<int64_t>& perm) const {
34+
const auto& shape = input.get_shape();
35+
const auto* input_data = input.as<T>();
36+
37+
if (!input_data || input_data->empty()) {
38+
throw std::runtime_error("Input tensor is empty or invalid");
39+
}
40+
41+
std::vector<size_t> new_dims;
42+
new_dims.reserve(shape.dims());
43+
for (const auto& axis : perm) {
44+
new_dims.push_back(shape[static_cast<size_t>(axis)]);
45+
}
46+
Shape new_shape(new_dims);
47+
48+
std::vector<size_t> input_strides(shape.dims());
49+
size_t stride = 1;
50+
for (size_t dim = shape.dims(); dim-- > 0;) {
51+
input_strides[dim] = stride;
52+
stride *= shape[dim];
53+
}
54+
55+
std::vector<size_t> output_strides(new_shape.dims());
56+
stride = 1;
57+
for (size_t dim = new_shape.dims(); dim-- > 0;) {
58+
output_strides[dim] = stride;
59+
stride *= new_shape[dim];
60+
}
61+
62+
std::vector<T> output_values(input_data->size());
63+
64+
if (shape.dims() == 2) {
65+
const size_t rows = shape[0];
66+
const size_t cols = shape[1];
67+
68+
if (perm[0] == 1 && perm[1] == 0) {
69+
for (size_t i = 0; i < rows; ++i) {
70+
for (size_t j = 0; j < cols; ++j) {
71+
const size_t new_index = j * rows + i;
72+
if (new_index >= output_values.size()) {
73+
throw std::runtime_error(
74+
"Index out of bounds during 2D transposition");
75+
}
76+
output_values[new_index] = (*input_data)[i * cols + j];
77+
}
78+
}
79+
} else {
80+
for (size_t i = 0; i < input_data->size(); ++i) {
81+
size_t old_index = i;
82+
size_t new_index = 0;
83+
84+
for (size_t dim = 0; dim < perm.size(); ++dim) {
85+
const auto axis = static_cast<size_t>(perm[dim]);
86+
const size_t coord = (old_index / input_strides[axis]) % shape[axis];
87+
new_index += coord * output_strides[dim];
88+
}
89+
90+
if (new_index >= output_values.size()) {
91+
throw std::runtime_error(
92+
"Index out of bounds during 2D transposition");
93+
}
94+
output_values[new_index] = (*input_data)[i];
95+
}
96+
}
97+
} else {
98+
for (size_t i = 0; i < input_data->size(); ++i) {
99+
size_t old_index = i;
100+
size_t new_index = 0;
101+
102+
for (size_t dim = 0; dim < perm.size(); ++dim) {
103+
const auto axis = static_cast<size_t>(perm[dim]);
104+
const size_t coord = (old_index / input_strides[axis]) % shape[axis];
105+
new_index += coord * output_strides[dim];
106+
}
107+
108+
if (new_index >= output_values.size()) {
109+
throw std::runtime_error("Index out of bounds during transposition");
110+
}
111+
output_values[new_index] = (*input_data)[i];
112+
}
113+
}
114+
115+
output = make_tensor(output_values, new_shape);
116+
}
117+
118+
void TransposeLayer::validate_perm(const Shape& input_shape,
119+
const std::vector<int64_t>& perm) {
120+
if (perm.size() != input_shape.dims()) {
121+
throw std::invalid_argument("Permutation size must match input dimensions");
122+
}
123+
124+
std::vector<bool> used_axes(input_shape.dims(), false);
125+
for (const auto& axis : perm) {
126+
if (axis < 0 || static_cast<size_t>(axis) >= input_shape.dims()) {
127+
throw std::invalid_argument("Invalid axis in permutation");
128+
}
129+
if (used_axes[static_cast<size_t>(axis)]) {
130+
throw std::invalid_argument("Duplicate axis in permutation");
131+
}
132+
used_axes[static_cast<size_t>(axis)] = true;
133+
}
134+
}
135+
136+
template void TransposeLayer::transpose_impl<float>(
137+
const Tensor&, Tensor&, const std::vector<int64_t>&) const;
138+
template void TransposeLayer::transpose_impl<int>(
139+
const Tensor&, Tensor&, const std::vector<int64_t>&) const;
140+
141+
} // namespace it_lab_ai

‎test/single_layer/test_transposelayer.cpp‎

Lines changed: 199 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,199 @@
1+
#include <cstdint>
2+
#include <iostream>
3+
#include <vector>
4+
5+
#include "gtest/gtest.h"
6+
#include "layers/Tensor.hpp"
7+
#include "layers/TransposeLayer.hpp"
8+
9+
using namespace it_lab_ai;
10+
11+
TEST(TransposeLayerTest, EmptyTensor) {
12+
Tensor input = make_tensor<float>({}, {0});
13+
TransposeLayer layer;
14+
Tensor output;
15+
16+
EXPECT_THROW(layer.run(input, output), std::runtime_error);
17+
}
18+
19+
TEST(TransposeLayerTest, IdentityTranspose) {
20+
Tensor input = make_tensor<float>({1, 2, 3, 4}, {2, 2});
21+
TransposeLayer layer({0, 1});
22+
Tensor output;
23+
24+
layer.run(input, output);
25+
26+
ASSERT_EQ(output.get_shape(), Shape({2, 2}));
27+
EXPECT_FLOAT_EQ(output.get<float>({0, 0}), 1.0f);
28+
EXPECT_FLOAT_EQ(output.get<float>({0, 1}), 2.0f);
29+
EXPECT_FLOAT_EQ(output.get<float>({1, 0}), 3.0f);
30+
EXPECT_FLOAT_EQ(output.get<float>({1, 1}), 4.0f);
31+
}
32+
33+
TEST(TransposeLayerTest, VectorTranspose) {
34+
Tensor input = make_tensor<float>({1, 2, 3, 4}, {4});
35+
TransposeLayer layer({0});
36+
Tensor output;
37+
38+
layer.run(input, output);
39+
40+
ASSERT_EQ(output.get_shape(), Shape({4}));
41+
EXPECT_FLOAT_EQ(output.get<float>({0}), 1.0f);
42+
EXPECT_FLOAT_EQ(output.get<float>({1}), 2.0f);
43+
EXPECT_FLOAT_EQ(output.get<float>({2}), 3.0f);
44+
EXPECT_FLOAT_EQ(output.get<float>({3}), 4.0f);
45+
}
46+
47+
TEST(TransposeLayerTest, InvalidPermutationSize) {
48+
Tensor input = make_tensor<float>({1, 2, 3, 4}, {2, 2});
49+
TransposeLayer layer({0});
50+
Tensor output;
51+
52+
EXPECT_THROW(layer.run(input, output), std::invalid_argument);
53+
}
54+
55+
TEST(TransposeLayerTest, DuplicateAxes) {
56+
Tensor input = make_tensor<float>({1, 2, 3, 4}, {2, 2});
57+
TransposeLayer layer({0, 0});
58+
Tensor output;
59+
60+
EXPECT_THROW(layer.run(input, output), std::invalid_argument);
61+
}
62+
63+
TEST(TransposeLayerTest, NegativeAxis) {
64+
Tensor input = make_tensor<float>({1, 2, 3, 4}, {2, 2});
65+
TransposeLayer layer({0, -1});
66+
Tensor output;
67+
68+
EXPECT_THROW(layer.run(input, output), std::invalid_argument);
69+
}
70+
71+
TEST(TransposeLayerTest, LargeAxis) {
72+
Tensor input = make_tensor<float>({1, 2, 3, 4}, {2, 2});
73+
TransposeLayer layer({0, 2});
74+
Tensor output;
75+
76+
EXPECT_THROW(layer.run(input, output), std::invalid_argument);
77+
}
78+
79+
TEST(TransposeLayerTest, 4DTensorTranspose) {
80+
std::vector<float> data(16);
81+
std::iota(data.begin(), data.end(), 1.0f);
82+
Tensor input = make_tensor<float>(data, {2, 2, 2, 2});
83+
TransposeLayer layer({3, 1, 0, 2});
84+
Tensor output;
85+
86+
layer.run(input, output);
87+
88+
ASSERT_EQ(output.get_shape(), Shape({2, 2, 2, 2}));
89+
EXPECT_FLOAT_EQ(output.get<float>({0, 0, 0, 0}), 1.0f);
90+
EXPECT_FLOAT_EQ(output.get<float>({1, 0, 0, 0}), 2.0f);
91+
EXPECT_FLOAT_EQ(output.get<float>({0, 1, 0, 0}), 5.0f);
92+
EXPECT_FLOAT_EQ(output.get<float>({1, 1, 0, 0}), 6.0f);
93+
EXPECT_FLOAT_EQ(output.get<float>({0, 0, 0, 1}), 3.0f);
94+
EXPECT_FLOAT_EQ(output.get<float>({1, 0, 0, 1}), 4.0f);
95+
EXPECT_FLOAT_EQ(output.get<float>({0, 1, 0, 1}), 7.0f);
96+
EXPECT_FLOAT_EQ(output.get<float>({1, 1, 0, 1}), 8.0f);
97+
EXPECT_FLOAT_EQ(output.get<float>({0, 0, 1, 0}), 9.0f);
98+
EXPECT_FLOAT_EQ(output.get<float>({1, 0, 1, 0}), 10.0f);
99+
EXPECT_FLOAT_EQ(output.get<float>({0, 1, 1, 0}), 13.0f);
100+
EXPECT_FLOAT_EQ(output.get<float>({1, 1, 1, 0}), 14.0f);
101+
EXPECT_FLOAT_EQ(output.get<float>({0, 0, 1, 1}), 11.0f);
102+
EXPECT_FLOAT_EQ(output.get<float>({1, 0, 1, 1}), 12.0f);
103+
EXPECT_FLOAT_EQ(output.get<float>({0, 1, 1, 1}), 15.0f);
104+
EXPECT_FLOAT_EQ(output.get<float>({1, 1, 1, 1}), 16.0f);
105+
}
106+
107+
TEST(TransposeLayerTest, MatrixTranspose) {
108+
Tensor input = make_tensor<float>({1, 2, 3, 4}, {2, 2});
109+
TransposeLayer layer({1, 0});
110+
Tensor output;
111+
112+
layer.run(input, output);
113+
114+
ASSERT_EQ(output.get_shape(), Shape({2, 2}));
115+
EXPECT_FLOAT_EQ(output.get<float>({0, 0}), 1.0f);
116+
EXPECT_FLOAT_EQ(output.get<float>({0, 1}), 3.0f);
117+
EXPECT_FLOAT_EQ(output.get<float>({1, 0}), 2.0f);
118+
EXPECT_FLOAT_EQ(output.get<float>({1, 1}), 4.0f);
119+
}
120+
121+
TEST(TransposeLayerTest, 3DTensor) {
122+
std::vector<float> data(24);
123+
std::iota(data.begin(), data.end(), 1.0f);
124+
Tensor input = make_tensor<float>(data, {2, 3, 4});
125+
TransposeLayer layer({2, 0, 1});
126+
Tensor output;
127+
128+
layer.run(input, output);
129+
130+
ASSERT_EQ(output.get_shape(), Shape({4, 2, 3}));
131+
EXPECT_FLOAT_EQ(output.get<float>({0, 0, 0}), 1.0f);
132+
EXPECT_FLOAT_EQ(output.get<float>({3, 1, 2}), 24.0f);
133+
}
134+
135+
TEST(TransposeLayerTest, IntTensor) {
136+
Tensor input = make_tensor<int>({1, 2, 3, 4}, {2, 2});
137+
TransposeLayer layer({1, 0});
138+
Tensor output;
139+
layer.run(input, output);
140+
EXPECT_EQ(output.get<int>({0, 0}), 1);
141+
EXPECT_EQ(output.get<int>({1, 0}), 2);
142+
EXPECT_EQ(output.get<int>({0, 1}), 3);
143+
EXPECT_EQ(output.get<int>({1, 1}), 4);
144+
}
145+
146+
TEST(TransposeLayerTest, 1DDefaultPermutationIsNoOp) {
147+
Tensor input = make_tensor<float>({1.0f, 2.0f, 3.0f, 4.0f}, {4});
148+
149+
TransposeLayer layer;
150+
Tensor output;
151+
152+
EXPECT_NO_THROW(layer.run(input, output));
153+
154+
EXPECT_EQ(output.get_shape(), Shape({4}));
155+
EXPECT_FLOAT_EQ(output.get<float>({0}), 1.0f);
156+
EXPECT_FLOAT_EQ(output.get<float>({1}), 2.0f);
157+
EXPECT_FLOAT_EQ(output.get<float>({2}), 3.0f);
158+
EXPECT_FLOAT_EQ(output.get<float>({3}), 4.0f);
159+
}
160+
161+
TEST(TransposeLayerTest, MultipleRunsWithDifferentRanks) {
162+
TransposeLayer layer({1, 0});
163+
164+
{
165+
Tensor input = make_tensor<float>({1, 2, 3, 4}, {2, 2});
166+
Tensor output;
167+
layer.run(input, output);
168+
EXPECT_EQ(output.get_shape(), Shape({2, 2}));
169+
EXPECT_FLOAT_EQ(output.get<float>({0, 0}), 1.0f);
170+
EXPECT_FLOAT_EQ(output.get<float>({1, 0}), 2.0f);
171+
}
172+
173+
{
174+
Tensor input = make_tensor<float>({1, 2, 3, 4, 5, 6}, {3, 2});
175+
Tensor output;
176+
layer.run(input, output);
177+
EXPECT_EQ(output.get_shape(), Shape({2, 3}));
178+
}
179+
180+
{
181+
Tensor input = make_tensor<float>({1, 2, 3, 4, 5, 6, 7, 8}, {2, 2, 2});
182+
Tensor output;
183+
EXPECT_THROW(layer.run(input, output), std::invalid_argument);
184+
}
185+
}
186+
187+
TEST(TransposeLayerTest, ExplicitPermutationWithDifferentRanks) {
188+
TransposeLayer layer({1, 0});
189+
190+
Tensor input2D = make_tensor<float>({1, 2, 3, 4}, {2, 2});
191+
Tensor output2D;
192+
layer.run(input2D, output2D);
193+
EXPECT_EQ(output2D.get_shape(), Shape({2, 2}));
194+
EXPECT_FLOAT_EQ(output2D.get<float>({1, 0}), 2.0f);
195+
196+
Tensor input3D = make_tensor<float>({1, 2, 3, 4, 5, 6, 7, 8}, {2, 2, 2});
197+
Tensor output3D;
198+
EXPECT_THROW(layer.run(input3D, output3D), std::invalid_argument);
199+
}

0 commit comments

Comments
 (0)