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