Function compute_softmax_node

cactus/graph/graph_ops_nn.cpp:560–581

}

void compute_softmax_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_buffer = get_input(node, 0, nodes, node_index_map);
    const auto& shape = input_buffer.shape;

    if (shape.size() < 2) {
        throw std::runtime_error("Softmax operation requires at least 2D tensor, got " +
                                std::to_string(shape.size()) + "D tensor");
    }

    if (input_buffer.precision != Precision::FP16) {
        throw std::runtime_error("Softmax operation only supports FP16 precision");
    }

    size_t batch_size = 1;
    for (size_t i = 0; i < shape.size() - 1; i++) {
        batch_size *= shape[i];
    }
    size_t vocab_size = shape[shape.size() - 1];

    cactus_softmax_f16(input_buffer.data_as<__fp16>(), node.output_buffer.data_as<__fp16>(),
                      batch_size, 1, vocab_size);
}

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