Skip to content

Commit 499c1e7

Browse files
authored
ReduceSum Layer (#187)
1 parent 7509f2f commit 499c1e7

File tree

4 files changed

+530
-8
lines changed

4 files changed

+530
-8
lines changed

include/layers/ReduceLayer.hpp

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,35 @@
1+
#pragma once
2+
#include <cstdint>
3+
#include <vector>
4+
5+
#include "layers/Layer.hpp"
6+
#include "layers/Tensor.hpp"
7+
8+
namespace it_lab_ai {
9+
10+
class ReduceLayer : public Layer {
11+
public:
12+
enum class Operation : uint8_t { kSum, kMean, kMult, kMax, kMin };
13+
14+
ReduceLayer(Operation op, int64_t keepdims = 0);
15+
explicit ReduceLayer(int64_t keepdims = 0)
16+
: ReduceLayer(Operation::kSum, keepdims) {}
17+
void run(const Tensor& input, Tensor& output) override;
18+
void run(const Tensor& input, const Tensor& axes, Tensor& output);
19+
20+
static std::string get_name() { return "ReduceLayer"; }
21+
22+
private:
23+
Operation op_;
24+
int64_t keepdims_;
25+
static void normalize_axes(const Shape& input_shape,
26+
std::vector<int64_t>& axes);
27+
Shape calculate_output_shape(const Shape& input_shape,
28+
const std::vector<int64_t>& axes) const;
29+
30+
template <typename T>
31+
void compute(const Tensor& input, const Shape& output_shape,
32+
const std::vector<int64_t>& axes, Tensor& output) const;
33+
};
34+
35+
} // namespace it_lab_ai

include/layers/Shape.hpp

Lines changed: 7 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -39,17 +39,16 @@ class Shape {
3939
}
4040
size_t dims() const noexcept { return dims_.size(); }
4141
size_t get_index(const std::vector<size_t>& coords) const;
42-
friend std::ostream& operator<<(std::ostream& os, const Shape& shape);
43-
bool operator==(const Shape& other) const noexcept {
44-
if (dims_.size() != other.dims_.size()) {
45-
return false;
42+
bool operator==(const Shape& other) const {
43+
if (dims_.size() != other.dims_.size()) return false;
44+
for (size_t i = 0; i < dims_.size(); ++i) {
45+
if (dims_[i] != other.dims_[i]) return false;
4646
}
47-
return std::equal(dims_.begin(), dims_.end(), other.dims_.begin());
47+
return true;
4848
}
4949

50-
bool operator!=(const Shape& other) const noexcept {
51-
return !(*this == other);
52-
}
50+
bool operator!=(const Shape& other) const { return !(*this == other); }
51+
friend std::ostream& operator<<(std::ostream& os, const Shape& shape);
5352

5453
private:
5554
std::vector<size_t> dims_;

src/layers/ReduceLayer.cpp

Lines changed: 213 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,213 @@
1+
#include "layers/ReduceLayer.hpp"
2+
3+
#include <algorithm>
4+
#include <limits>
5+
#include <numeric>
6+
7+
namespace it_lab_ai {
8+
9+
ReduceLayer::ReduceLayer(Operation op, int64_t keepdims)
10+
: op_(op), keepdims_(keepdims) {}
11+
12+
void ReduceLayer::normalize_axes(const Shape& input_shape,
13+
std::vector<int64_t>& axes) {
14+
const auto rank = static_cast<int64_t>(input_shape.dims());
15+
16+
if (rank == 0) {
17+
if (!axes.empty()) {
18+
throw std::runtime_error("ReduceLayer: Axis specified for scalar input");
19+
}
20+
return;
21+
}
22+
23+
if (axes.empty()) {
24+
axes.resize(rank);
25+
std::iota(axes.begin(), axes.end(), 0);
26+
return;
27+
}
28+
29+
for (auto& axis : axes) {
30+
if (axis < -rank || axis >= rank) {
31+
throw std::runtime_error(
32+
"ReduceLayer: Axis out of range. Valid range is [-" +
33+
std::to_string(rank) + ", " + std::to_string(rank - 1) + "]");
34+
}
35+
36+
if (axis < 0) {
37+
axis += rank;
38+
}
39+
}
40+
41+
std::sort(axes.begin(), axes.end());
42+
axes.erase(std::unique(axes.begin(), axes.end()), axes.end());
43+
}
44+
45+
Shape ReduceLayer::calculate_output_shape(
46+
const Shape& input_shape, const std::vector<int64_t>& axes) const {
47+
if (input_shape.dims() == 0) {
48+
return Shape({});
49+
}
50+
51+
std::vector<size_t> new_dims;
52+
53+
if (keepdims_) {
54+
new_dims.resize(input_shape.dims(), 1);
55+
for (int64_t i = 0; i < static_cast<int64_t>(input_shape.dims()); ++i) {
56+
bool is_axis = std::find(axes.begin(), axes.end(), i) != axes.end();
57+
if (!is_axis) {
58+
new_dims[i] = input_shape[i];
59+
}
60+
}
61+
} else {
62+
for (int64_t i = 0; i < static_cast<int64_t>(input_shape.dims()); ++i) {
63+
bool is_axis = std::find(axes.begin(), axes.end(), i) != axes.end();
64+
if (!is_axis) {
65+
new_dims.push_back(input_shape[i]);
66+
}
67+
}
68+
if (new_dims.empty()) {
69+
new_dims.push_back(1);
70+
}
71+
}
72+
73+
return Shape(new_dims);
74+
}
75+
76+
template <typename T>
77+
void ReduceLayer::compute(const Tensor& input, const Shape& output_shape,
78+
const std::vector<int64_t>& axes,
79+
Tensor& output) const {
80+
const auto& input_data = *input.as<T>();
81+
std::vector<T> output_data(output_shape.count());
82+
std::vector<size_t> counts(output_shape.count(), 0);
83+
84+
switch (op_) {
85+
case Operation::kSum:
86+
case Operation::kMean:
87+
std::fill(output_data.begin(), output_data.end(), T(0));
88+
break;
89+
case Operation::kMult:
90+
std::fill(output_data.begin(), output_data.end(), T(1));
91+
break;
92+
case Operation::kMax:
93+
std::fill(output_data.begin(), output_data.end(),
94+
std::numeric_limits<T>::lowest());
95+
break;
96+
case Operation::kMin:
97+
std::fill(output_data.begin(), output_data.end(),
98+
std::numeric_limits<T>::max());
99+
break;
100+
}
101+
102+
const auto& input_shape = input.get_shape();
103+
const auto input_rank = static_cast<int64_t>(input_shape.dims());
104+
105+
std::vector<size_t> in_coords(input_rank, 0);
106+
for (size_t in_idx = 0; in_idx < input_data.size(); ++in_idx) {
107+
std::vector<size_t> out_coords;
108+
if (keepdims_) {
109+
out_coords.resize(input_rank, 0);
110+
for (int64_t i = 0; i < input_rank; ++i) {
111+
if (std::find(axes.begin(), axes.end(), i) == axes.end()) {
112+
out_coords[i] = in_coords[i];
113+
}
114+
}
115+
} else {
116+
for (int64_t i = 0; i < input_rank; ++i) {
117+
if (std::find(axes.begin(), axes.end(), i) == axes.end()) {
118+
out_coords.push_back(in_coords[i]);
119+
}
120+
}
121+
}
122+
123+
size_t out_idx = 0;
124+
size_t stride = 1;
125+
for (size_t i = out_coords.size(); i-- > 0;) {
126+
out_idx += out_coords[i] * stride;
127+
stride *= output_shape[i];
128+
}
129+
130+
switch (op_) {
131+
case Operation::kSum:
132+
case Operation::kMean:
133+
output_data[out_idx] += input_data[in_idx];
134+
counts[out_idx]++;
135+
break;
136+
case Operation::kMult:
137+
output_data[out_idx] *= input_data[in_idx];
138+
break;
139+
case Operation::kMax:
140+
if (input_data[in_idx] > output_data[out_idx]) {
141+
output_data[out_idx] = input_data[in_idx];
142+
}
143+
break;
144+
case Operation::kMin:
145+
if (input_data[in_idx] < output_data[out_idx]) {
146+
output_data[out_idx] = input_data[in_idx];
147+
}
148+
break;
149+
}
150+
151+
for (int64_t i = input_rank; i-- > 0;) {
152+
++in_coords[i];
153+
if (in_coords[i] < input_shape[i]) break;
154+
in_coords[i] = 0;
155+
}
156+
}
157+
158+
if (op_ == Operation::kMean) {
159+
for (size_t i = 0; i < output_data.size(); ++i) {
160+
if (counts[i] != 0) {
161+
output_data[i] /= static_cast<T>(counts[i]);
162+
}
163+
}
164+
}
165+
166+
output = make_tensor(output_data, output_shape);
167+
}
168+
169+
template void ReduceLayer::compute<float>(const Tensor&, const Shape&,
170+
const std::vector<int64_t>&,
171+
Tensor&) const;
172+
template void ReduceLayer::compute<int>(const Tensor&, const Shape&,
173+
const std::vector<int64_t>&,
174+
Tensor&) const;
175+
176+
void ReduceLayer::run(const Tensor& input, Tensor& output) {
177+
run(input, Tensor(), output);
178+
}
179+
180+
void ReduceLayer::run(const Tensor& input, const Tensor& axes, Tensor& output) {
181+
if (input.get_shape().count() == 0) {
182+
output = make_tensor<float>({0.0F}, {});
183+
return;
184+
}
185+
186+
std::vector<int64_t> axes_indices;
187+
if (axes.get_shape().dims() > 0) {
188+
if (axes.get_type() == Type::kInt) {
189+
const auto* axes_data = axes.as<int>();
190+
axes_indices.assign(axes_data->begin(), axes_data->end());
191+
} else {
192+
throw std::runtime_error("ReduceLayer: Axes tensor must be of type int");
193+
}
194+
}
195+
196+
normalize_axes(input.get_shape(), axes_indices);
197+
Shape output_shape = calculate_output_shape(input.get_shape(), axes_indices);
198+
199+
switch (input.get_type()) {
200+
case Type::kFloat:
201+
compute<float>(input, output_shape, axes_indices, output);
202+
break;
203+
case Type::kInt:
204+
compute<int>(input, output_shape, axes_indices, output);
205+
break;
206+
default:
207+
throw std::runtime_error(
208+
"ReduceLayer: Unsupported input tensor type. Only float and int are "
209+
"supported");
210+
}
211+
}
212+
213+
} // namespace it_lab_ai

0 commit comments

Comments
 (0)