Skip to content

Commit baf492d

Browse files
Dobiasdclaude
Dobiasd
and
claude
committed
Refactor adaptive_pooling, einsum_dense, GQA for readability
Split deeply-nested apply_impl bodies into named helpers: - adaptive_pooling_3d_layer: extracted pool_window() and adapt_range() (also folded the in_size==1 special case into adapt_range), plus output_shape_for() to keep the apply_impl loop body shallow. - einsum_dense_layer: extracted derive_char_sizes(), contract_one(), add_bias(), and decode_index/encode_offset helpers. - group_query_attention_layer: split into project(), attention_distribution(), compute_attention(), apply_sigmoid_gate(), and output_projection(). No functional change — all existing tests still pass. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
1 parent 41afc71 commit baf492d

3 files changed

Lines changed: 267 additions & 230 deletions

File tree

‎include/fdeep/layers/adaptive_pooling_3d_layer.hpp‎

Lines changed: 58 additions & 50 deletions
Original file line numberDiff line numberDiff line change
@@ -40,73 +40,81 @@ namespace internal {
4040
const std::size_t out_w_;
4141
const adaptive_pooling_kind kind_;
4242

43-
static std::size_t adapt_start(std::size_t i, std::size_t in_size, std::size_t out_size)
43+
struct range {
44+
std::size_t start;
45+
std::size_t end;
46+
};
47+
48+
// Adaptive pooling maps output index i to input range
49+
// [floor(i * in / out), ceil((i+1) * in / out)). For length-1 input
50+
// dimensions (which can occur for promoted 1D/2D tensors) the range
51+
// collapses to [0, 1).
52+
static range adapt_range(std::size_t i, std::size_t in_size, std::size_t out_size)
4453
{
45-
return static_cast<std::size_t>(std::floor(
54+
if (in_size == 1)
55+
return { 0, 1 };
56+
const auto start = static_cast<std::size_t>(std::floor(
4657
static_cast<double>(i * in_size) / static_cast<double>(out_size)));
58+
const auto end = static_cast<std::size_t>(std::ceil(
59+
static_cast<double>((i + 1) * in_size) / static_cast<double>(out_size)));
60+
return { start, end };
4761
}
4862

49-
static std::size_t adapt_end(std::size_t i, std::size_t in_size, std::size_t out_size)
63+
static tensor_shape output_shape_for(const tensor_shape& in_shape,
64+
std::size_t out_d4, std::size_t out_h, std::size_t out_w)
5065
{
51-
return static_cast<std::size_t>(std::ceil(
52-
static_cast<double>((i + 1) * in_size) / static_cast<double>(out_size)));
66+
const std::size_t depth = in_shape.depth_;
67+
switch (in_shape.rank()) {
68+
case 2:
69+
return tensor_shape(out_w, depth);
70+
case 3:
71+
return tensor_shape(out_h, out_w, depth);
72+
default:
73+
return tensor_shape(out_d4, out_h, out_w, depth);
74+
}
75+
}
76+
77+
float_type pool_window(const tensor& input,
78+
range d, range h, range w, std::size_t z) const
79+
{
80+
const bool is_max = kind_ == adaptive_pooling_kind::max;
81+
float_type acc = is_max
82+
? std::numeric_limits<float_type>::lowest()
83+
: float_type(0);
84+
std::size_t count = 0;
85+
for (std::size_t di = d.start; di < d.end; ++di) {
86+
for (std::size_t yi = h.start; yi < h.end; ++yi) {
87+
for (std::size_t xi = w.start; xi < w.end; ++xi) {
88+
const float_type v = input.get_ignore_rank(tensor_pos(0, di, yi, xi, z));
89+
acc = is_max ? std::max(acc, v) : acc + v;
90+
++count;
91+
}
92+
}
93+
}
94+
if (!is_max && count > 0)
95+
acc /= static_cast<float_type>(count);
96+
return acc;
5397
}
5498

5599
tensors apply_impl(const tensors& inputs) const override
56100
{
57101
const auto& input = single_tensor_from_tensors(inputs);
58102
const auto& sh = input.shape();
59-
const std::size_t in_d4 = sh.size_dim_4_;
60-
const std::size_t in_h = sh.height_;
61-
const std::size_t in_w = sh.width_;
62-
const std::size_t depth = sh.depth_;
63-
const std::size_t out_d4 = out_d4_ == 0 ? in_d4 : out_d4_;
64-
const std::size_t out_h = out_h_ == 0 ? in_h : out_h_;
103+
const std::size_t out_d4 = out_d4_ == 0 ? sh.size_dim_4_ : out_d4_;
104+
const std::size_t out_h = out_h_ == 0 ? sh.height_ : out_h_;
65105
const std::size_t out_w = out_w_;
66106

67-
tensor_shape out_shape(out_d4, out_h, out_w, depth);
68-
// Match input rank when possible (so 1D/2D inputs produce 1D/2D outputs).
69-
if (sh.rank() <= 3) {
70-
if (sh.rank() == 2)
71-
out_shape = tensor_shape(out_w, depth);
72-
else
73-
out_shape = tensor_shape(out_h, out_w, depth);
74-
} else if (sh.rank() == 4) {
75-
out_shape = tensor_shape(out_d4, out_h, out_w, depth);
76-
}
77-
78-
tensor out(out_shape, float_type(0));
107+
tensor out(output_shape_for(sh, out_d4, out_h, out_w), float_type(0));
79108

80109
for (std::size_t od = 0; od < out_d4; ++od) {
81-
const std::size_t d_start = in_d4 == 1 ? 0 : adapt_start(od, in_d4, out_d4);
82-
const std::size_t d_end = in_d4 == 1 ? 1 : adapt_end(od, in_d4, out_d4);
110+
const range d = adapt_range(od, sh.size_dim_4_, out_d4);
83111
for (std::size_t oy = 0; oy < out_h; ++oy) {
84-
const std::size_t y_start = in_h == 1 ? 0 : adapt_start(oy, in_h, out_h);
85-
const std::size_t y_end = in_h == 1 ? 1 : adapt_end(oy, in_h, out_h);
112+
const range h = adapt_range(oy, sh.height_, out_h);
86113
for (std::size_t ox = 0; ox < out_w; ++ox) {
87-
const std::size_t x_start = adapt_start(ox, in_w, out_w);
88-
const std::size_t x_end = adapt_end(ox, in_w, out_w);
89-
for (std::size_t z = 0; z < depth; ++z) {
90-
float_type acc = kind_ == adaptive_pooling_kind::max
91-
? std::numeric_limits<float_type>::lowest()
92-
: float_type(0);
93-
std::size_t count = 0;
94-
for (std::size_t d = d_start; d < d_end; ++d) {
95-
for (std::size_t y = y_start; y < y_end; ++y) {
96-
for (std::size_t x = x_start; x < x_end; ++x) {
97-
const float_type v = input.get_ignore_rank(
98-
tensor_pos(0, d, y, x, z));
99-
if (kind_ == adaptive_pooling_kind::max)
100-
acc = std::max(acc, v);
101-
else
102-
acc += v;
103-
++count;
104-
}
105-
}
106-
}
107-
if (kind_ == adaptive_pooling_kind::avg && count > 0)
108-
acc /= static_cast<float_type>(count);
109-
out.set_ignore_rank(tensor_pos(0, od, oy, ox, z), acc);
114+
const range w = adapt_range(ox, sh.width_, out_w);
115+
for (std::size_t z = 0; z < sh.depth_; ++z) {
116+
out.set_ignore_rank(tensor_pos(0, od, oy, ox, z),
117+
pool_window(input, d, h, w, z));
110118
}
111119
}
112120
}

0 commit comments

Comments
 (0)