Function compute_stft_node

cactus/graph/graph_ops_nn.cpp:1817–1837

}

void compute_stft_node(GraphNode& node, const std::vector<std::unique_ptr<GraphNode>>& nodes,
                                 const std::unordered_map<size_t, size_t>& node_index_map) {
    const auto& X = get_input(node, 0, nodes, node_index_map);
    const auto& W = get_input(node, 1, nodes, node_index_map);
    auto& Y = node.output_buffer;

    const size_t N = X.shape[0];
    const size_t C_in = X.shape[1];
    const size_t L = X.shape[2];
    const size_t C_out = W.shape[0];
    const size_t K = W.shape[2];
    const size_t stride = node.params.stride;
    const size_t num_fft_bins = node.params.num_fft_bins;

    if (X.precision != Precision::FP16 || W.precision != Precision::FP16) {
        throw std::runtime_error("stft only supports FP16");
    }

    cactus_stft_f16(X.data_as<__fp16>(), W.data_as<__fp16>(),
                            Y.data_as<__fp16>(), N, L, C_in, C_out, K, stride, num_fft_bins);
}

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