Function compute_stats_pool_node

cactus/graph/graph_ops_nn.cpp:2276–2305

}

void compute_stats_pool_node(GraphNode& node, const std::vector<std::unique_ptr<GraphNode>>& nodes,
                              const std::unordered_map<size_t, size_t>& node_index_map) {
    const auto& input = get_input(node, 0, nodes, node_index_map);
    const __fp16* src = input.data_as<__fp16>();
    __fp16* dst = node.output_buffer.data_as<__fp16>();

    size_t batch = input.shape[0];
    size_t total_per_batch = input.total_size / batch;
    size_t T = input.shape.back();
    size_t features = total_per_batch / T;

    for (size_t b = 0; b < batch; ++b) {
        const __fp16* batch_src = src + b * total_per_batch;
        __fp16* batch_dst = dst + b * features * 2;

        for (size_t f = 0; f < features; ++f) {
            float sum = 0.0f, sum_sq = 0.0f;
            for (size_t t = 0; t < T; ++t) {
                float v = static_cast<float>(batch_src[f * T + t]);
                sum += v;
                sum_sq += v * v;
            }
            float mean = sum / static_cast<float>(T);
            float var = T > 1 ? (sum_sq - static_cast<float>(T) * mean * mean) / static_cast<float>(T - 1) : 0.0f;
            float std_val = sqrtf(fmaxf(var, 0.0f));
            batch_dst[f] = static_cast<__fp16>(mean);
            batch_dst[features + f] = static_cast<__fp16>(std_val);
        }
    }
}

void compute_weighted_stats_pool_node(GraphNode& node, const std::vector<std::unique_ptr<GraphNode>>& nodes,
                                       const std::unordered_map<size_t, size_t>& node_index_map) {