/*
 * Copyright (c) 2020-2024, NVIDIA CORPORATION. All rights reserved.
 *
 * This source code and/or documentation ("Licensed Deliverables") are
 * subject to NVIDIA intellectual property rights under U.S. and
 * international Copyright laws.
 */

#pragma once

#include <cugraph-ops/cuda/stream.hpp>
#include <cugraph-ops/graph/format.hpp>
#include <cugraph-ops/operators/common.hpp>

#include <cstdint>

namespace cugraph::ops {

/**
 * @brief Computes the forward pass for a "GAT-like" multi-head attention layer
 *        (mha_gat) operating on a graph in a node-to-node
 *        reduction (n2n).
 *
 * @note If `graph.map_csc_to_coo` are set, then we use this as an edge index to obtain
 * the index of an edge feature given the index of an edge.
 * @note With an arbitrary edge feature index, the size of `edge_feat` is
 * arbitrary.
 * @note We always assume `dim_head = dim_node / params.num_heads`.
 *
 * @param[out] out             the output embeddings after attention.
 *                             [on device] [dim = `graph.n_dst_nodes x dim_out`].
 *                             For each node, the features associated with
 *                             the different heads are concatenated if
 *                             params.concat_heads = true (`dim_out = dim_node`),
 *                             and averaged otherwise (`dim_out = dim_head`).
 * @param[inout] sm_scores     the output softmax scores for the backward pass.
 *                             [on device] [dim = `2 x params.num_heads x n_edges`]
 * @param[in]  node_feat       the embeddings of all input nodes. [on device]
 *                             [dim = `graph.n_dst_nodes x dim_node`]
 *                             For each node, the features associated with
 *                             the different heads are expected to be concatenated.
 * @param[in]  attn_weights    the weights applied to in/out feat [on device]
 *                             [dim = `dim_node + dim_node`]
 * @param[in]  graph           the input graph.
 * @param[in]  params          the multi-head attention parameters
 * @param[in]  dim_node        dimension of the input embeddings.
 *                             It must be a multiple of `params.num_heads`.
 * @param[in]  stream          cuda stream
 * @{
 */
void mha_gat_n2n_fwd(float* out,
                     float* sm_scores,
                     const float* node_feat,
                     const float* attn_weights,
                     const csc_s32_t& graph,
                     const mha_params& params,
                     int dim_node,
                     const cuda::stream& stream);
void mha_gat_n2n_fwd(float* out,
                     float* sm_scores,
                     const float* node_feat,
                     const float* attn_weights,
                     const csc_s64_t& graph,
                     const mha_params& params,
                     int dim_node,
                     const cuda::stream& stream);
void mha_gat_n2n_fwd(__half* out,
                     float* sm_scores,
                     const __half* node_feat,
                     const __half* attn_weights,
                     const csc_s32_t& graph,
                     const mha_params& params,
                     int dim_node,
                     const cuda::stream& stream);
void mha_gat_n2n_fwd(__half* out,
                     float* sm_scores,
                     const __half* node_feat,
                     const __half* attn_weights,
                     const csc_s64_t& graph,
                     const mha_params& params,
                     int dim_node,
                     const cuda::stream& stream);
void mha_gat_n2n_fwd(__nv_bfloat16* out,
                     float* sm_scores,
                     const __nv_bfloat16* node_feat,
                     const __nv_bfloat16* attn_weights,
                     const csc_s32_t& graph,
                     const mha_params& params,
                     int dim_node,
                     const cuda::stream& stream);
void mha_gat_n2n_fwd(__nv_bfloat16* out,
                     float* sm_scores,
                     const __nv_bfloat16* node_feat,
                     const __nv_bfloat16* attn_weights,
                     const csc_s64_t& graph,
                     const mha_params& params,
                     int dim_node,
                     const cuda::stream& stream);
/** @} */

/**
 * @brief Computes the forward pass for a "GAT-like" multi-head attention layer
 *        (mha_gat) operating on a graph in a node-to-node reduction (n2n)
 *        but using edge features, too, for computing the dot product (efeat).
 *
 * @note If `graph.map_csc_to_coo` are set, then we use this as an edge index to obtain
 * the index of an edge feature given the index of an edge.
 * @note With an arbitrary edge feature index, the size of `edge_feat` is
 * arbitrary.
 * @note We always assume `dim_head = dim_node / params.num_heads`.
 *
 * @param[out] out             the output embeddings after attention.
 *                             [on device] [dim = `graph.n_dst_nodes x dim_out`].
 *                             For each node, the features associated with
 *                             the different heads are concatenated if
 *                             params.concat_heads = true (`dim_out = dim_node`),
 *                             and averaged otherwise (`dim_out = dim_head`).
 * @param[inout] sm_scores     the output softmax scores for the backward pass.
 *                             [on device] [dim = `2 x params.num_heads x n_edges`]
 * @param[in]  node_feat       the embeddings of all input nodes. [on device]
 *                             [dim = `graph.n_dst_nodes x dim_node`]
 *                             For each node, the features associated with
 *                             the different heads are expected to be concatenated.
 * @param[in]  edge_feat       the embeddings of all edges [on device]
 *                             [dim = `n_edges x dim_edge`]
 *                             For each edge, the features associated with the different
 *                             heads are expected to be concatenated.
 * @param[in]  attn_weights    the weights applied to in/out feat and edge features [on device]
 *                             [dim = `dim_node + dim_node + dim_edge`]
 * @param[in]  graph           the input graph.
 * @param[in]  params          the multi-head attention parameters
 * @param[in]  dim_node        dimension of the input embeddings.
 *                             It must be a multiple of `params.num_heads`.
 * @param[in]  dim_edge        dimension of the edge embeddings.
 *                             It must be a multiple of `params.num_heads`.
 * @param[in]  stream          cuda stream
 * @{
 */
void mha_gat_n2n_efeat_fwd(float* out,
                           float* sm_scores,
                           const float* node_feat,
                           const float* edge_feat,
                           const float* attn_weights,
                           const csc_s32_t& graph,
                           const mha_params& params,
                           int dim_node,
                           int dim_edge,
                           const cuda::stream& stream);
void mha_gat_n2n_efeat_fwd(float* out,
                           float* sm_scores,
                           const float* node_feat,
                           const float* edge_feat,
                           const float* attn_weights,
                           const csc_s64_t& graph,
                           const mha_params& params,
                           int dim_node,
                           int dim_edge,
                           const cuda::stream& stream);
void mha_gat_n2n_efeat_fwd(__nv_bfloat16* out,
                           float* sm_scores,
                           const __nv_bfloat16* node_feat,
                           const __nv_bfloat16* edge_feat,
                           const __nv_bfloat16* attn_weights,
                           const csc_s32_t& graph,
                           const mha_params& params,
                           int dim_node,
                           int dim_edge,
                           const cuda::stream& stream);
void mha_gat_n2n_efeat_fwd(__nv_bfloat16* out,
                           float* sm_scores,
                           const __nv_bfloat16* node_feat,
                           const __nv_bfloat16* edge_feat,
                           const __nv_bfloat16* attn_weights,
                           const csc_s64_t& graph,
                           const mha_params& params,
                           int dim_node,
                           int dim_edge,
                           const cuda::stream& stream);
void mha_gat_n2n_efeat_fwd(__half* out,
                           float* sm_scores,
                           const __half* node_feat,
                           const __half* edge_feat,
                           const __half* attn_weights,
                           const csc_s32_t& graph,
                           const mha_params& params,
                           int dim_node,
                           int dim_edge,
                           const cuda::stream& stream);
void mha_gat_n2n_efeat_fwd(__half* out,
                           float* sm_scores,
                           const __half* node_feat,
                           const __half* edge_feat,
                           const __half* attn_weights,
                           const csc_s64_t& graph,
                           const mha_params& params,
                           int dim_node,
                           int dim_edge,
                           const cuda::stream& stream);
/** @} */

/**
 * @brief Computes the backward pass for a "GAT-like" multi-head attention layer
 *        (mha_gat) operating on a graph in a node-to-node reduction (n2n).
 *
 * @note If `graph.map_csc_to_coo` are set, then we use this as an edge index to obtain
 * the index of an edge feature given the index of an edge.
 * @note With an arbitrary edge feature index, the size of `edge_feat` is
 * arbitrary.
 * @note We always assume `dim_head = dim_node / params.num_heads`.
 * @note `grad_node_feat` and `grad_weights` may be set to `nullptr` if not needed.
 *
 * @param[out] grad_node_feat       the gradients on input embeddings of all nodes
 *                                  [on device] [dim = `graph.n_dst_nodes x dim_node`].
 * @param[out] grad_weights         the gradients on all attention weights [on device]
 *                                  [dim = `dim_node + dim_node`]
 * @param[inout] grad_wspace_feat   workspace for deterministic feature-gradient reduction,
 *                                  if passed as none, non-deterministic path is chosen
 *                                  [on device]
 *                                  [dim = `graph.n_indices * dim_node]
 * @param[inout] grad_wspace_weight workspace for deterministic gradient reduction,
 *                                  if passed as none, non-deterministic path is chosen
 *                                  [on device]
 *                                  [dim = graph.n_dst_nodes * dim_weight`]
 * @param[inout] grad_sm_scores     buffer for the gradient on softmax scores.
 *                                  [on device] [dim = `2 x params.num_heads x n_edges`]
 * @param[in]  grad_out             the gradient on the output of the forward pass
 *                                  [on device] [dim = `graph.n_dst_nodes x dim_out`]
 *                                  For each node, the gradients associated with
 *                                  the different heads are assumed to be
 *                                  separate if params.concat_heads = true (`dim_out = dim_node`),
 *                                  and shared otherwise (`dim_out = dim_head`).
 * @param[in]  node_feat            the embeddings of all input nodes. [on device]
 *                                  [dim = `graph.n_dst_nodes x dim_node`]
 *                                  For each node, the features associated with
 *                                  the different heads are expected to be concatenated.
 * @param[in]  attn_weights         the weights applied to in/out feat [on device]
 *                                  [dim = `dim_node + dim_node`]
 * @param[inout]  sm_scores         the output softmax scores from the forward pass.
 *                                  [on device] [dim = `2 x params.num_heads x n_edges`]
 * @param[in]  graph                the input graph.
 * @param[in]  params               the multi-head attention parameters
 * @param[in]  dim_node             dimension of the input embeddings.
 *                                  It must be a multiple of `params.num_heads`.
 * @param[in]  stream               cuda stream
 * @{
 */
void mha_gat_n2n_bwd(float* grad_node_feat,
                     float* grad_weights,
                     float* grad_wspace_feat,
                     float* grad_wspace_weight,
                     float* grad_sm_scores,
                     const float* grad_out,
                     const float* node_feat,
                     const float* attn_weights,
                     const float* sm_scores,
                     const csc_s32_t& graph,
                     const mha_params& params,
                     int dim_node,
                     const cuda::stream& stream);
void mha_gat_n2n_bwd(float* grad_node_feat,
                     float* grad_weights,
                     float* grad_wspace_feat,
                     float* grad_wspace_weight,
                     float* grad_sm_scores,
                     const float* grad_out,
                     const float* node_feat,
                     const float* attn_weights,
                     const float* sm_scores,
                     const csc_s64_t& graph,
                     const mha_params& params,
                     int dim_node,
                     const cuda::stream& stream);
void mha_gat_n2n_bwd(__half* grad_node_feat,
                     __half* grad_weights,
                     __half* grad_wspace_feat,
                     __half* grad_wspace_weight,
                     float* grad_sm_scores,
                     const __half* grad_out,
                     const __half* node_feat,
                     const __half* attn_weights,
                     const float* sm_scores,
                     const csc_s32_t& graph,
                     const mha_params& params,
                     int dim_node,
                     const cuda::stream& stream);
void mha_gat_n2n_bwd(__half* grad_node_feat,
                     __half* grad_weights,
                     __half* grad_wspace_feat,
                     __half* grad_wspace_weight,
                     float* grad_sm_scores,
                     const __half* grad_out,
                     const __half* node_feat,
                     const __half* attn_weights,
                     const float* sm_scores,
                     const csc_s64_t& graph,
                     const mha_params& params,
                     int dim_node,
                     const cuda::stream& stream);
void mha_gat_n2n_bwd(__nv_bfloat16* grad_node_feat,
                     __nv_bfloat16* grad_weights,
                     __nv_bfloat16* grad_wspace_feat,
                     __nv_bfloat16* grad_wspace_weight,
                     float* grad_sm_scores,
                     const __nv_bfloat16* grad_out,
                     const __nv_bfloat16* node_feat,
                     const __nv_bfloat16* attn_weights,
                     const float* sm_scores,
                     const csc_s32_t& graph,
                     const mha_params& params,
                     int dim_node,
                     const cuda::stream& stream);
void mha_gat_n2n_bwd(__nv_bfloat16* grad_node_feat,
                     __nv_bfloat16* grad_weights,
                     __nv_bfloat16* grad_wspace_feat,
                     __nv_bfloat16* grad_wspace_weight,
                     float* grad_sm_scores,
                     const __nv_bfloat16* grad_out,
                     const __nv_bfloat16* node_feat,
                     const __nv_bfloat16* attn_weights,
                     const float* sm_scores,
                     const csc_s64_t& graph,
                     const mha_params& params,
                     int dim_node,
                     const cuda::stream& stream);
void mha_gat_n2n_bwd(float* grad_node_feat,
                     float* grad_weights,
                     float* grad_wspace_feat,
                     float* grad_wspace_weight,
                     float* grad_sm_scores,
                     const __half* grad_out,
                     const __half* node_feat,
                     const __half* attn_weights,
                     const float* sm_scores,
                     const csc_s32_t& graph,
                     const mha_params& params,
                     int dim_node,
                     const cuda::stream& stream);
void mha_gat_n2n_bwd(float* grad_node_feat,
                     float* grad_weights,
                     float* grad_wspace_feat,
                     float* grad_wspace_weight,
                     float* grad_sm_scores,
                     const __half* grad_out,
                     const __half* node_feat,
                     const __half* attn_weights,
                     const float* sm_scores,
                     const csc_s64_t& graph,
                     const mha_params& params,
                     int dim_node,
                     const cuda::stream& stream);
void mha_gat_n2n_bwd(float* grad_node_feat,
                     float* grad_weights,
                     float* grad_wspace_feat,
                     float* grad_wspace_weight,
                     float* grad_sm_scores,
                     const __nv_bfloat16* grad_out,
                     const __nv_bfloat16* node_feat,
                     const __nv_bfloat16* attn_weights,
                     const float* sm_scores,
                     const csc_s32_t& graph,
                     const mha_params& params,
                     int dim_node,
                     const cuda::stream& stream);
void mha_gat_n2n_bwd(float* grad_node_feat,
                     float* grad_weights,
                     float* grad_wspace_feat,
                     float* grad_wspace_weight,
                     float* grad_sm_scores,
                     const __nv_bfloat16* grad_out,
                     const __nv_bfloat16* node_feat,
                     const __nv_bfloat16* attn_weights,
                     const float* sm_scores,
                     const csc_s64_t& graph,
                     const mha_params& params,
                     int dim_node,
                     const cuda::stream& stream);
void mha_gat_n2n_bwd(__half* grad_node_feat,
                     float* grad_weights,
                     __half* grad_wspace_feat,
                     float* grad_wspace_weight,
                     float* grad_sm_scores,
                     const __half* grad_out,
                     const __half* node_feat,
                     const __half* attn_weights,
                     const float* sm_scores,
                     const csc_s32_t& graph,
                     const mha_params& params,
                     int dim_node,
                     const cuda::stream& stream);
void mha_gat_n2n_bwd(__half* grad_node_feat,
                     float* grad_weights,
                     __half* grad_wspace_feat,
                     float* grad_wspace_weight,
                     float* grad_sm_scores,
                     const __half* grad_out,
                     const __half* node_feat,
                     const __half* attn_weights,
                     const float* sm_scores,
                     const csc_s64_t& graph,
                     const mha_params& params,
                     int dim_node,
                     const cuda::stream& stream);
void mha_gat_n2n_bwd(__nv_bfloat16* grad_node_feat,
                     float* grad_weights,
                     __nv_bfloat16* grad_wspace_feat,
                     float* grad_wspace_weight,
                     float* grad_sm_scores,
                     const __nv_bfloat16* grad_out,
                     const __nv_bfloat16* node_feat,
                     const __nv_bfloat16* attn_weights,
                     const float* sm_scores,
                     const csc_s32_t& graph,
                     const mha_params& params,
                     int dim_node,
                     const cuda::stream& stream);
void mha_gat_n2n_bwd(__nv_bfloat16* grad_node_feat,
                     float* grad_weights,
                     __nv_bfloat16* grad_wspace_feat,
                     float* grad_wspace_weight,
                     float* grad_sm_scores,
                     const __nv_bfloat16* grad_out,
                     const __nv_bfloat16* node_feat,
                     const __nv_bfloat16* attn_weights,
                     const float* sm_scores,
                     const csc_s64_t& graph,
                     const mha_params& params,
                     int dim_node,
                     const cuda::stream& stream);

/** @} */

/**
 * @brief Computes the backward pass for a "GAT-like" multi-head attention layer
 *        (mha_gat) operating on a graph in a node-to-node reduction (n2n)
 *        but using edge features, too, for computing the dot product (efeat).
 *
 * @note If `graph.map_csc_to_coo` are set, then we use this as an edge index to obtain
 * the index of an edge feature given the index of an edge.
 * @note With an arbitrary edge feature index, the size of `edge_feat` is
 * arbitrary.
 * @note We always assume `dim_head = dim_node / params.num_heads`.
 * @note `grad_node_feat`, `grad_edge_feat` and `grad_weights` may be set to
 * `nullptr` if not needed.
 *
 * @param[out] grad_node_feat       the gradients on input embeddings of all nodes
 *                                  [on device] [dim = `graph.n_dst_nodes x dim_node`].
 * @param[out] grad_edge_feat       the gradient on embeddings of all edges [on device]
 *                                  [dim = `n_edges x dim_edge`]
 * @param[out] grad_weights         the gradients on all attention weights [on device]
 *                                  [dim = `dim_node + dim_node + dim_edge`]
 * @param[inout] grad_wspace_feat   workspace for deterministic feature-gradient reduction,
 *                                  if passed as none, non-deterministic path is chosen
 *                                  [on device]
 *                                  [dim = `graph.n_indices * dim_node]
 * @param[inout] grad_wspace_weight workspace for deterministic gradient reduction,
 *                                  if passed as none, non-deterministic path is chosen
 *                                  [on device]
 *                                  [dim = graph.n_dst_nodes * dim_weight`]
 * @param[inout] grad_sm_scores     the gradient on softmax scores.
 *                                  [on device] [dim = `2 x params.num_heads x n_edges`]
 * @param[in]  grad_out             the gradient on the output of the forward pass
 *                                  [on device] [dim = `graph.n_dst_nodes x dim_node`]
 * @param[in]  node_feat            the embeddings of all input nodes.
 *                                  [on device] [dim = `graph.n_dst_nodes x dim_out`]
 *                                  For each node, the gradients associated with
 *                                  the different heads are assumed to be
 *                                  separate if params.concat_heads = true (`dim_out = dim_node`),
 *                                  and shared otherwise (`dim_out = dim_head`).
 * @param[in]  edge_feat            the embeddings of all edges [on device]
 *                                  [dim = `n_edges x dim_edge`]
 *                                  For each edge, the features associated with the different
 *                                  heads are expected to be concatenated.
 * @param[in]  attn_weights         the weights applied to in/out feat [on device]
 *                                  [dim = `dim_node + dim_node + dim_edge`]
 * @param[inout]  sm_scores         the output softmax scores from the forward pass.
 *                                  [on device] [dim = `2 x params.num_heads x n_edges`]
 * @param[in]  graph                the input graph.
 * @param[in]  params               the multi-head attention parameters
 * @param[in]  dim_node             dimension of the input embeddings.
 *                                  It must be a multiple of `params.num_heads`.
 * @param[in]  dim_edge             dimension of the edge embeddings
 *                                  It must be a multiple of `params.num_heads`.
 * @param[in]  stream               cuda stream
 * @{
 */
void mha_gat_n2n_efeat_bwd(float* grad_node_feat,
                           float* grad_edge_feat,
                           float* grad_weights,
                           float* grad_wspace_feat,
                           float* grad_wspace_weight,
                           float* grad_sm_scores,
                           const float* grad_out,
                           const float* node_feat,
                           const float* edge_feat,
                           const float* attn_weights,
                           const float* sm_scores,
                           const csc_s32_t& graph,
                           const mha_params& params,
                           int dim_node,
                           int dim_edge,
                           const cuda::stream& stream);
void mha_gat_n2n_efeat_bwd(float* grad_node_feat,
                           float* grad_edge_feat,
                           float* grad_weights,
                           float* grad_wspace_feat,
                           float* grad_wspace_weight,
                           float* grad_sm_scores,
                           const float* grad_out,
                           const float* node_feat,
                           const float* edge_feat,
                           const float* attn_weights,
                           const float* sm_scores,
                           const csc_s64_t& graph,
                           const mha_params& params,
                           int dim_node,
                           int dim_edge,
                           const cuda::stream& stream);

void mha_gat_n2n_efeat_bwd(__nv_bfloat16* grad_node_feat,
                           __nv_bfloat16* grad_edge_feat,
                           __nv_bfloat16* grad_weights,
                           __nv_bfloat16* grad_wspace_feat,
                           __nv_bfloat16* grad_wspace_weight,
                           float* grad_sm_scores,
                           const __nv_bfloat16* grad_out,
                           const __nv_bfloat16* node_feat,
                           const __nv_bfloat16* edge_feat,
                           const __nv_bfloat16* attn_weights,
                           const float* sm_scores,
                           const csc_s32_t& graph,
                           const mha_params& params,
                           int dim_node,
                           int dim_edge,
                           const cuda::stream& stream);
void mha_gat_n2n_efeat_bwd(__nv_bfloat16* grad_node_feat,
                           __nv_bfloat16* grad_edge_feat,
                           __nv_bfloat16* grad_weights,
                           __nv_bfloat16* grad_wspace_feat,
                           __nv_bfloat16* grad_wspace_weight,
                           float* grad_sm_scores,
                           const __nv_bfloat16* grad_out,
                           const __nv_bfloat16* node_feat,
                           const __nv_bfloat16* edge_feat,
                           const __nv_bfloat16* attn_weights,
                           const float* sm_scores,
                           const csc_s64_t& graph,
                           const mha_params& params,
                           int dim_node,
                           int dim_edge,
                           const cuda::stream& stream);

void mha_gat_n2n_efeat_bwd(__half* grad_node_feat,
                           __half* grad_edge_feat,
                           __half* grad_weights,
                           __half* grad_wspace_feat,
                           __half* grad_wspace_weight,
                           float* grad_sm_scores,
                           const __half* grad_out,
                           const __half* node_feat,
                           const __half* edge_feat,
                           const __half* attn_weights,
                           const float* sm_scores,
                           const csc_s32_t& graph,
                           const mha_params& params,
                           int dim_node,
                           int dim_edge,
                           const cuda::stream& stream);
void mha_gat_n2n_efeat_bwd(__half* grad_node_feat,
                           __half* grad_edge_feat,
                           __half* grad_weights,
                           __half* grad_wspace_feat,
                           __half* grad_wspace_weight,
                           float* grad_sm_scores,
                           const __half* grad_out,
                           const __half* node_feat,
                           const __half* edge_feat,
                           const __half* attn_weights,
                           const float* sm_scores,
                           const csc_s64_t& graph,
                           const mha_params& params,
                           int dim_node,
                           int dim_edge,
                           const cuda::stream& stream);

void mha_gat_n2n_efeat_bwd(float* grad_node_feat,
                           float* grad_edge_feat,
                           float* grad_weights,
                           float* grad_wspace_feat,
                           float* grad_wspace_weight,
                           float* grad_sm_scores,
                           const __nv_bfloat16* grad_out,
                           const __nv_bfloat16* node_feat,
                           const __nv_bfloat16* edge_feat,
                           const __nv_bfloat16* attn_weights,
                           const float* sm_scores,
                           const csc_s32_t& graph,
                           const mha_params& params,
                           int dim_node,
                           int dim_edge,
                           const cuda::stream& stream);
void mha_gat_n2n_efeat_bwd(float* grad_node_feat,
                           float* grad_edge_feat,
                           float* grad_weights,
                           float* grad_wspace_feat,
                           float* grad_wspace_weight,
                           float* grad_sm_scores,
                           const __nv_bfloat16* grad_out,
                           const __nv_bfloat16* node_feat,
                           const __nv_bfloat16* edge_feat,
                           const __nv_bfloat16* attn_weights,
                           const float* sm_scores,
                           const csc_s64_t& graph,
                           const mha_params& params,
                           int dim_node,
                           int dim_edge,
                           const cuda::stream& stream);

void mha_gat_n2n_efeat_bwd(float* grad_node_feat,
                           float* grad_edge_feat,
                           float* grad_weights,
                           float* grad_wspace_feat,
                           float* grad_wspace_weight,
                           float* grad_sm_scores,
                           const __half* grad_out,
                           const __half* node_feat,
                           const __half* edge_feat,
                           const __half* attn_weights,
                           const float* sm_scores,
                           const csc_s32_t& graph,
                           const mha_params& params,
                           int dim_node,
                           int dim_edge,
                           const cuda::stream& stream);
void mha_gat_n2n_efeat_bwd(float* grad_node_feat,
                           float* grad_edge_feat,
                           float* grad_weights,
                           float* grad_wspace_feat,
                           float* grad_wspace_weight,
                           float* grad_sm_scores,
                           const __half* grad_out,
                           const __half* node_feat,
                           const __half* edge_feat,
                           const __half* attn_weights,
                           const float* sm_scores,
                           const csc_s64_t& graph,
                           const mha_params& params,
                           int dim_node,
                           int dim_edge,
                           const cuda::stream& stream);

void mha_gat_n2n_efeat_bwd(__nv_bfloat16* grad_node_feat,
                           __nv_bfloat16* grad_edge_feat,
                           float* grad_weights,
                           __nv_bfloat16* grad_wspace_feat,
                           float* grad_wspace_weight,
                           float* grad_sm_scores,
                           const __nv_bfloat16* grad_out,
                           const __nv_bfloat16* node_feat,
                           const __nv_bfloat16* edge_feat,
                           const __nv_bfloat16* attn_weights,
                           const float* sm_scores,
                           const csc_s32_t& graph,
                           const mha_params& params,
                           int dim_node,
                           int dim_edge,
                           const cuda::stream& stream);
void mha_gat_n2n_efeat_bwd(__nv_bfloat16* grad_node_feat,
                           __nv_bfloat16* grad_edge_feat,
                           float* grad_weights,
                           __nv_bfloat16* grad_wspace_feat,
                           float* grad_wspace_weight,
                           float* grad_sm_scores,
                           const __nv_bfloat16* grad_out,
                           const __nv_bfloat16* node_feat,
                           const __nv_bfloat16* edge_feat,
                           const __nv_bfloat16* attn_weights,
                           const float* sm_scores,
                           const csc_s64_t& graph,
                           const mha_params& params,
                           int dim_node,
                           int dim_edge,
                           const cuda::stream& stream);

void mha_gat_n2n_efeat_bwd(__half* grad_node_feat,
                           __half* grad_edge_feat,
                           float* grad_weights,
                           __half* grad_wspace_feat,
                           float* grad_wspace_weight,
                           float* grad_sm_scores,
                           const __half* grad_out,
                           const __half* node_feat,
                           const __half* edge_feat,
                           const __half* attn_weights,
                           const float* sm_scores,
                           const csc_s32_t& graph,
                           const mha_params& params,
                           int dim_node,
                           int dim_edge,
                           const cuda::stream& stream);
void mha_gat_n2n_efeat_bwd(__half* grad_node_feat,
                           __half* grad_edge_feat,
                           float* grad_weights,
                           __half* grad_wspace_feat,
                           float* grad_wspace_weight,
                           float* grad_sm_scores,
                           const __half* grad_out,
                           const __half* node_feat,
                           const __half* edge_feat,
                           const __half* attn_weights,
                           const float* sm_scores,
                           const csc_s64_t& graph,
                           const mha_params& params,
                           int dim_node,
                           int dim_edge,
                           const cuda::stream& stream);
/** @} */

/**
 * @brief Computes the forward pass for a "GAT-like" multi-head attention layer
 *        (mha_gat) operating on bipartite graphs in a node-to-node
 *        reduction (n2n).
 *
 * @note If `graph.map_csc_to_coo` are set, then we use this as an edge index to obtain
 * the index of an edge feature given the index of an edge.
 * @note We always assume `dim_head = dim_node / params.num_heads`.
 *
 * @param[out] out             the output embeddings after attention.
 *                             [on device] [dim = `graph.n_dst_nodes x dim_out`].
 *                             For each node, the features associated with
 *                             the different heads are concatenated if
 *                             params.concat_heads = true (`dim_out = dim_node`),
 *                             and averaged otherwise (`dim_out = dim_head`).
 * @param[inout] sm_scores     the output softmax scores for the backward pass.
 *                             [on device] [dim = `2 x params.num_heads x graph.n_indices`]
 * @param[in]  src_feat        the embeddings of all source nodes. [on device]
 *                             [dim = `graph.n_src_nodes x dim_node`]
 *                             For each node, the features associated with
 *                             the different heads are expected to be concatenated.
 * @param[in]  dst_feat        the embeddings of all destination nodes. [on device]
 *                             [dim = `graph.n_dst_nodes x dim_node`]
 *                             For each node, the features associated with
 *                             the different heads are expected to be concatenated.
 * @param[in]  attn_weights    the weights applied to in/out feat [on device]
 *                             [dim = `dim_node + dim_node`]
 * @param[in]  graph           the input graph (`bipartite_csc`).
 * @param[in]  params          the multi-head attention parameters
 * @param[in]  dim_node        dimension of the input embeddings.
 *                             It must be a multiple of `params.num_heads`.
 * @param[in]  stream          cuda stream
 * @{
 */
void mha_gat_n2n_fwd(float* out,
                     float* sm_scores,
                     const float* src_feat,
                     const float* dst_feat,
                     const float* attn_weights,
                     const bipartite_csc_s32_t& graph,
                     const mha_params& params,
                     int dim_node,
                     const cuda::stream& stream);
void mha_gat_n2n_fwd(float* out,
                     float* sm_scores,
                     const float* src_feat,
                     const float* dst_feat,
                     const float* attn_weights,
                     const bipartite_csc_s64_t& graph,
                     const mha_params& params,
                     int dim_node,
                     const cuda::stream& stream);
void mha_gat_n2n_fwd(__half* out,
                     float* sm_scores,
                     const __half* src_feat,
                     const __half* dst_feat,
                     const __half* attn_weights,
                     const bipartite_csc_s32_t& graph,
                     const mha_params& params,
                     int dim_node,
                     const cuda::stream& stream);
void mha_gat_n2n_fwd(__half* out,
                     float* sm_scores,
                     const __half* src_feat,
                     const __half* dst_feat,
                     const __half* attn_weights,
                     const bipartite_csc_s64_t& graph,
                     const mha_params& params,
                     int dim_node,
                     const cuda::stream& stream);
void mha_gat_n2n_fwd(__nv_bfloat16* out,
                     float* sm_scores,
                     const __nv_bfloat16* src_feat,
                     const __nv_bfloat16* dst_feat,
                     const __nv_bfloat16* attn_weights,
                     const bipartite_csc_s32_t& graph,
                     const mha_params& params,
                     int dim_node,
                     const cuda::stream& stream);
void mha_gat_n2n_fwd(__nv_bfloat16* out,
                     float* sm_scores,
                     const __nv_bfloat16* src_feat,
                     const __nv_bfloat16* dst_feat,
                     const __nv_bfloat16* attn_weights,
                     const bipartite_csc_s64_t& graph,
                     const mha_params& params,
                     int dim_node,
                     const cuda::stream& stream);
/** @} */

/**
 * @brief Computes the forward pass for a "GAT-like" multi-head attention layer
 *        (mha_gat) operating on bipartite graphs in a node-to-node reduction (n2n)
 *        but using edge features, too, for computing the dot product (efeat).
 *
 * @note If `graph.map_csc_to_coo` are set, then we use this as an edge index to obtain
 * the index of an edge feature given the index of an edge.
 * @note We always assume `dim_head = dim_node / params.num_heads`.
 *
 * @param[out] out             the output embeddings after attention.
 *                             [on device] [dim = `graph.n_dst_nodes x dim_out`].
 *                             For each node, the features associated with
 *                             the different heads are concatenated if
 *                             params.concat_heads = true (`dim_out = dim_node`),
 *                             and averaged otherwise (`dim_out = dim_head`).
 * @param[inout] sm_scores     the output softmax scores for the backward pass.
 *                             [on device] [dim = `2 x params.num_heads x graph.n_indices`]
 * @param[in]  src_feat        the embeddings of all source nodes. [on device]
 *                             [dim = `graph.n_src_nodes x dim_node`]
 *                             For each node, the features associated with
 *                             the different heads are expected to be concatenated.
 * @param[in]  dst_feat        the embeddings of all destination nodes. [on device]
 *                             [dim = `graph.n_dst_nodes x dim_node`]
 *                             For each node, the features associated with
 *                             the different heads are expected to be concatenated.
 * @param[in]  edge_feat       the embeddings of all edges [on device]
 *                             [dim = `graph.n_indices x dim_edge`]
 *                             For each edge, the features associated with the different
 *                             heads are expected to be concatenated.
 * @param[in]  attn_weights    the weights applied to in/out feat and edge features [on device]
 *                             [dim = `dim_node + dim_node + dim_edge`]
 * @param[in]  graph           the input graph (`bipartite_csc`).
 * @param[in]  params          the multi-head attention parameters
 * @param[in]  dim_node        dimension of the input embeddings.
 *                             It must be a multiple of `params.num_heads`.
 * @param[in]  dim_edge        dimension of the edge embeddings.
 *                             It must be a multiple of `params.num_heads`.
 * @param[in]  stream          cuda stream
 * @{
 */
void mha_gat_n2n_efeat_fwd(float* out,
                           float* sm_scores,
                           const float* src_feat,
                           const float* dst_feat,
                           const float* edge_feat,
                           const float* attn_weights,
                           const bipartite_csc_s32_t& graph,
                           const mha_params& params,
                           int dim_node,
                           int dim_edge,
                           const cuda::stream& stream);
void mha_gat_n2n_efeat_fwd(float* out,
                           float* sm_scores,
                           const float* src_feat,
                           const float* dst_feat,
                           const float* edge_feat,
                           const float* attn_weights,
                           const bipartite_csc_s64_t& graph,
                           const mha_params& params,
                           int dim_node,
                           int dim_edge,
                           const cuda::stream& stream);
void mha_gat_n2n_efeat_fwd(__nv_bfloat16* out,
                           float* sm_scores,
                           const __nv_bfloat16* src_feat,
                           const __nv_bfloat16* dst_feat,
                           const __nv_bfloat16* edge_feat,
                           const __nv_bfloat16* attn_weights,
                           const bipartite_csc_s32_t& graph,
                           const mha_params& params,
                           int dim_node,
                           int dim_edge,
                           const cuda::stream& stream);
void mha_gat_n2n_efeat_fwd(__nv_bfloat16* out,
                           float* sm_scores,
                           const __nv_bfloat16* src_feat,
                           const __nv_bfloat16* dst_feat,
                           const __nv_bfloat16* edge_feat,
                           const __nv_bfloat16* attn_weights,
                           const bipartite_csc_s64_t& graph,
                           const mha_params& params,
                           int dim_node,
                           int dim_edge,
                           const cuda::stream& stream);
void mha_gat_n2n_efeat_fwd(__half* out,
                           float* sm_scores,
                           const __half* src_feat,
                           const __half* dst_feat,
                           const __half* edge_feat,
                           const __half* attn_weights,
                           const bipartite_csc_s32_t& graph,
                           const mha_params& params,
                           int dim_node,
                           int dim_edge,
                           const cuda::stream& stream);
void mha_gat_n2n_efeat_fwd(__half* out,
                           float* sm_scores,
                           const __half* src_feat,
                           const __half* dst_feat,
                           const __half* edge_feat,
                           const __half* attn_weights,
                           const bipartite_csc_s64_t& graph,
                           const mha_params& params,
                           int dim_node,
                           int dim_edge,
                           const cuda::stream& stream);
/** @} */

/**
 * @brief Computes the backward pass for a "GAT-like" multi-head attention layer
 *        (mha_gat) operating on bipartite graphs in a node-to-node reduction (n2n)
 *
 * @note If `graph.map_csc_to_coo` are set, then we use this as an edge index to obtain
 * the index of an edge feature given the index of an edge.
 * @note We always assume `dim_head = dim_node / params.num_heads`.
 * @note `grad_src_feat`, `grad_dst_feat` and `grad_weights` may be set to `nullptr` if not needed.
 *
 * @param[out] grad_src_feat        the gradients on input embeddings of all source nodes
 *                                  [on device] [dim = `graph.n_src_nodes x dim_node`].
 * @param[out] grad_dst_feat        the gradients on input embeddings of all destination nodes
 *                                  [on device] [dim = `graph.n_dst_nodes x dim_node`].
 * @param[out] grad_weights         the gradients on all attention weights [on device]
 *                                  [dim = `dim_node + dim_node`]
 * @param[inout] grad_wspace_feat   workspace for deterministic feature-gradient reduction,
 *                                  if passed as none, non-deterministic path is chosen
 *                                  [on device]
 *                                  [dim = `graph.n_indices * dim_node]
 * @param[inout] grad_wspace_weight workspace for deterministic gradient reduction,
 *                                  if passed as none, non-deterministic path is chosen
 *                                  [on device]
 *                                  [dim = graph.n_dst_nodes * dim_weight`]
 * @param[inout] grad_sm_scores     buffer for the gradient on softmax scores.
 *                                  [on device] [dim = `2 x params.num_heads x graph.n_indices`]
 * @param[in]  grad_out             the gradient on the output of the forward pass
 *                                  [on device] [dim = `graph.n_dst_nodes x dim_out`]
 *                                  For each node, the gradients associated with
 *                                  the different heads are assumed to be
 *                                  separate if params.concat_heads = true (`dim_out = dim_node`),
 *                                  and shared otherwise (`dim_out = dim_head`).
 * @param[in]  src_feat             the embeddings of all source nodes. [on device]
 *                                  [dim = `graph.n_src_nodes x dim_node`]
 *                                  For each node, the features associated with
 *                                  the different heads are expected to be concatenated.
 * @param[in]  dst_feat             the embeddings of all destination nodes. [on device]
 *                                  [dim = `graph.n_dst_nodes x dim_node`]
 *                                  For each node, the features associated with
 *                                  the different heads are expected to be concatenated.
 * @param[in]  attn_weights         the weights applied to in/out feat [on device]
 *                                  [dim = `dim_node + dim_node`]
 * @param[inout] sm_scores          the softmax scores from the forward pass.
 *                                  [on device] [dim = `2 x params.num_heads x graph.n_indices`]
 * @param[in]  graph                the input graph (`bipartite_csc`).
 * @param[in]  params               the multi-head attention parameters
 * @param[in]  dim_node             dimension of the input embeddings.
 *                                  It must be a multiple of `params.num_heads`.
 * @param[in]  stream               cuda stream
 * @{
 */
void mha_gat_n2n_bwd(float* grad_src_feat,
                     float* grad_dst_feat,
                     float* grad_weights,
                     float* grad_wspace_feat,
                     float* grad_wspace_weight,
                     float* grad_sm_scores,
                     const float* grad_out,
                     const float* src_feat,
                     const float* dst_feat,
                     const float* attn_weights,
                     const float* sm_scores,
                     const bipartite_csc_s32_t& graph,
                     const mha_params& params,
                     int dim_node,
                     const cuda::stream& stream);
void mha_gat_n2n_bwd(float* grad_src_feat,
                     float* grad_dst_feat,
                     float* grad_weights,
                     float* grad_wspace_feat,
                     float* grad_wspace_weight,
                     float* grad_sm_scores,
                     const float* grad_out,
                     const float* src_feat,
                     const float* dst_feat,
                     const float* attn_weights,
                     const float* sm_scores,
                     const bipartite_csc_s64_t& graph,
                     const mha_params& params,
                     int dim_node,
                     const cuda::stream& stream);
void mha_gat_n2n_bwd(__half* grad_src_feat,
                     __half* grad_dst_feat,
                     __half* grad_weights,
                     __half* grad_wspace_feat,
                     __half* grad_wspace_weight,
                     float* grad_sm_scores,
                     const __half* grad_out,
                     const __half* src_feat,
                     const __half* dst_feat,
                     const __half* attn_weights,
                     const float* sm_scores,
                     const bipartite_csc_s32_t& graph,
                     const mha_params& params,
                     int dim_node,
                     const cuda::stream& stream);
void mha_gat_n2n_bwd(__half* grad_src_feat,
                     __half* grad_dst_feat,
                     __half* grad_weights,
                     __half* grad_wspace_feat,
                     __half* grad_wspace_weight,
                     float* grad_sm_scores,
                     const __half* grad_out,
                     const __half* src_feat,
                     const __half* dst_feat,
                     const __half* attn_weights,
                     const float* sm_scores,
                     const bipartite_csc_s64_t& graph,
                     const mha_params& params,
                     int dim_node,
                     const cuda::stream& stream);
void mha_gat_n2n_bwd(__nv_bfloat16* grad_src_feat,
                     __nv_bfloat16* grad_dst_feat,
                     __nv_bfloat16* grad_weights,
                     __nv_bfloat16* grad_wspace_feat,
                     __nv_bfloat16* grad_wspace_weight,
                     float* grad_sm_scores,
                     const __nv_bfloat16* grad_out,
                     const __nv_bfloat16* src_feat,
                     const __nv_bfloat16* dst_feat,
                     const __nv_bfloat16* attn_weights,
                     const float* sm_scores,
                     const bipartite_csc_s32_t& graph,
                     const mha_params& params,
                     int dim_node,
                     const cuda::stream& stream);
void mha_gat_n2n_bwd(__nv_bfloat16* grad_src_feat,
                     __nv_bfloat16* grad_dst_feat,
                     __nv_bfloat16* grad_weights,
                     __nv_bfloat16* grad_wspace_feat,
                     __nv_bfloat16* grad_wspace_weight,
                     float* grad_sm_scores,
                     const __nv_bfloat16* grad_out,
                     const __nv_bfloat16* src_feat,
                     const __nv_bfloat16* dst_feat,
                     const __nv_bfloat16* attn_weights,
                     const float* sm_scores,
                     const bipartite_csc_s64_t& graph,
                     const mha_params& params,
                     int dim_node,
                     const cuda::stream& stream);
void mha_gat_n2n_bwd(float* grad_src_feat,
                     float* grad_dst_feat,
                     float* grad_weights,
                     float* grad_wspace_feat,
                     float* grad_wspace_weight,
                     float* grad_sm_scores,
                     const __half* grad_out,
                     const __half* src_feat,
                     const __half* dst_feat,
                     const __half* attn_weights,
                     const float* sm_scores,
                     const bipartite_csc_s32_t& graph,
                     const mha_params& params,
                     int dim_node,
                     const cuda::stream& stream);
void mha_gat_n2n_bwd(float* grad_src_feat,
                     float* grad_dst_feat,
                     float* grad_weights,
                     float* grad_wspace_feat,
                     float* grad_wspace_weight,
                     float* grad_sm_scores,
                     const __half* grad_out,
                     const __half* src_feat,
                     const __half* dst_feat,
                     const __half* attn_weights,
                     const float* sm_scores,
                     const bipartite_csc_s64_t& graph,
                     const mha_params& params,
                     int dim_node,
                     const cuda::stream& stream);
void mha_gat_n2n_bwd(float* grad_src_feat,
                     float* grad_dst_feat,
                     float* grad_weights,
                     float* grad_wspace_feat,
                     float* grad_wspace_weight,
                     float* grad_sm_scores,
                     const __nv_bfloat16* grad_out,
                     const __nv_bfloat16* src_feat,
                     const __nv_bfloat16* dst_feat,
                     const __nv_bfloat16* attn_weights,
                     const float* sm_scores,
                     const bipartite_csc_s32_t& graph,
                     const mha_params& params,
                     int dim_node,
                     const cuda::stream& stream);
void mha_gat_n2n_bwd(float* grad_src_feat,
                     float* grad_dst_feat,
                     float* grad_weights,
                     float* grad_wspace_feat,
                     float* grad_wspace_weight,
                     float* grad_sm_scores,
                     const __nv_bfloat16* grad_out,
                     const __nv_bfloat16* src_feat,
                     const __nv_bfloat16* dst_feat,
                     const __nv_bfloat16* attn_weights,
                     const float* sm_scores,
                     const bipartite_csc_s64_t& graph,
                     const mha_params& params,
                     int dim_node,
                     const cuda::stream& stream);

void mha_gat_n2n_bwd(__half* grad_src_feat,
                     __half* grad_dst_feat,
                     float* grad_weights,
                     __half* grad_wspace_feat,
                     float* grad_wspace_weight,
                     float* grad_sm_scores,
                     const __half* grad_out,
                     const __half* src_feat,
                     const __half* dst_feat,
                     const __half* attn_weights,
                     const float* sm_scores,
                     const bipartite_csc_s32_t& graph,
                     const mha_params& params,
                     int dim_node,
                     const cuda::stream& stream);
void mha_gat_n2n_bwd(__half* grad_src_feat,
                     __half* grad_dst_feat,
                     float* grad_weights,
                     __half* grad_wspace_feat,
                     float* grad_wspace_weight,
                     float* grad_sm_scores,
                     const __half* grad_out,
                     const __half* src_feat,
                     const __half* dst_feat,
                     const __half* attn_weights,
                     const float* sm_scores,
                     const bipartite_csc_s64_t& graph,
                     const mha_params& params,
                     int dim_node,
                     const cuda::stream& stream);
void mha_gat_n2n_bwd(__nv_bfloat16* grad_src_feat,
                     __nv_bfloat16* grad_dst_feat,
                     float* grad_weights,
                     __nv_bfloat16* grad_wspace_feat,
                     float* grad_wspace_weight,
                     float* grad_sm_scores,
                     const __nv_bfloat16* grad_out,
                     const __nv_bfloat16* src_feat,
                     const __nv_bfloat16* dst_feat,
                     const __nv_bfloat16* attn_weights,
                     const float* sm_scores,
                     const bipartite_csc_s32_t& graph,
                     const mha_params& params,
                     int dim_node,
                     const cuda::stream& stream);
void mha_gat_n2n_bwd(__nv_bfloat16* grad_src_feat,
                     __nv_bfloat16* grad_dst_feat,
                     float* grad_weights,
                     __nv_bfloat16* grad_wspace_feat,
                     float* grad_wspace_weight,
                     float* grad_sm_scores,
                     const __nv_bfloat16* grad_out,
                     const __nv_bfloat16* src_feat,
                     const __nv_bfloat16* dst_feat,
                     const __nv_bfloat16* attn_weights,
                     const float* sm_scores,
                     const bipartite_csc_s64_t& graph,
                     const mha_params& params,
                     int dim_node,
                     const cuda::stream& stream);
/** @} */

/**
 * @brief Computes the backward pass for a "GAT-like" multi-head attention layer
 *        (mha_gat) operating on bipartite graphs in a node-to-node reduction (n2n)
 *        but using edge features, too, for computing the dot product (efeat).
 *
 * @note If `graph.map_csc_to_coo` are set, then we use this as an edge index to obtain
 * the index of an edge feature given the index of an edge.
 * @note We always assume `dim_head = dim_node / params.num_heads`.
 * @note `grad_node_feat`, `grad_edge_feat` and `grad_weights` may be set to
 * `nullptr` if not needed.
 *
 * @param[out] grad_src_feat        the gradients on input embeddings of all source nodes
 *                                  [on device] [dim = `graph.n_src_nodes x dim_node`].
 * @param[out] grad_dst_feat        the gradients on input embeddings of all destination nodes
 *                                  [on device] [dim = `graph.n_dst_nodes x dim_node`].
 * @param[out] grad_edge_feat       the gradient on embeddings of all edges [on device]
 *                                  [dim = `graph.n_indices x dim_edge`]
 * @param[out] grad_weights         the gradients on all attention weights [on device]
 *                                  [dim = `dim_node + dim_node + dim_edge`]
 * @param[inout] grad_wspace_feat   workspace for deterministic feature-gradient reduction,
 *                                  if passed as none, non-deterministic path is chosen
 *                                  [on device]
 *                                  [dim = `graph.n_indices * dim_node]
 * @param[inout] grad_wspace_weight workspace for deterministic gradient reduction,
 *                                  if passed as none, non-deterministic path is chosen
 *                                  [on device]
 *                                  [dim = graph.n_dst_nodes * dim_weight`]
 * @param[inout] grad_sm_scores     the gradient on softmax scores.
 *                                  [on device] [dim = `2 x params.num_heads x graph.n_indices`]
 * @param[in]  grad_out             the gradient on the output of the forward pass
 *                                  [on device] [dim = `graph.n_dst_nodes x dim_node`]
 * @param[in]  src_feat             the embeddings of all source nodes. [on device]
 *                                  [dim = `graph.n_src_nodes x dim_node`]
 *                                  For each node, the features associated with
 *                                  the different heads are expected to be concatenated.
 * @param[in]  dst_feat             the embeddings of all destination nodes. [on device]
 *                                  [dim = `graph.n_dst_nodes x dim_node`]
 *                                  For each node, the features associated with
 *                                  the different heads are expected to be concatenated.
 * @param[in]  edge_feat            the embeddings of all edges [on device]
 *                                  [dim = `graph.n_indices x dim_edge`]
 *                                  For each edge, the features associated with the different
 *                                  heads are expected to be concatenated.
 * @param[in]  attn_weights         the weights applied to in/out feat [on device]
 *                                  [dim = `dim_node + dim_node + dim_edge`]
 * @param[inout] sm_scores          the output softmax scores from the forward pass.
 *                                  [on device] [dim = `2 x params.num_heads x graph.n_indices`]
 * @param[in]  graph                the input graph (`bipartite_csc`).
 * @param[in]  params               the multi-head attention parameters
 * @param[in]  dim_node             dimension of the input embeddings.
 *                                  It must be a multiple of `params.num_heads`.
 * @param[in]  dim_edge             dimension of the edge embeddings
 *                                  It must be a multiple of `params.num_heads`.
 * @param[in]  stream               cuda stream
 * @{
 */
void mha_gat_n2n_efeat_bwd(float* grad_src_feat,
                           float* grad_dst_feat,
                           float* grad_edge_feat,
                           float* grad_weights,
                           float* grad_wspace_feat,
                           float* grad_wspace_weight,
                           float* grad_sm_scores,
                           const float* grad_out,
                           const float* src_feat,
                           const float* dst_feat,
                           const float* edge_feat,
                           const float* attn_weights,
                           const float* sm_scores,
                           const bipartite_csc_s32_t& graph,
                           const mha_params& params,
                           int dim_node,
                           int dim_edge,
                           const cuda::stream& stream);
void mha_gat_n2n_efeat_bwd(float* grad_src_feat,
                           float* grad_dst_feat,
                           float* grad_edge_feat,
                           float* grad_weights,
                           float* grad_wspace_feat,
                           float* grad_wspace_weight,
                           float* grad_sm_scores,
                           const float* grad_out,
                           const float* src_feat,
                           const float* dst_feat,
                           const float* edge_feat,
                           const float* attn_weights,
                           const float* sm_scores,
                           const bipartite_csc_s64_t& graph,
                           const mha_params& params,
                           int dim_node,
                           int dim_edge,
                           const cuda::stream& stream);

void mha_gat_n2n_efeat_bwd(__nv_bfloat16* grad_src_feat,
                           __nv_bfloat16* grad_dst_feat,
                           __nv_bfloat16* grad_edge_feat,
                           __nv_bfloat16* grad_weights,
                           __nv_bfloat16* grad_wspace_feat,
                           __nv_bfloat16* grad_wspace_weight,
                           float* grad_sm_scores,
                           const __nv_bfloat16* grad_out,
                           const __nv_bfloat16* src_feat,
                           const __nv_bfloat16* dst_feat,
                           const __nv_bfloat16* edge_feat,
                           const __nv_bfloat16* attn_weights,
                           const float* sm_scores,
                           const bipartite_csc_s32_t& graph,
                           const mha_params& params,
                           int dim_node,
                           int dim_edge,
                           const cuda::stream& stream);
void mha_gat_n2n_efeat_bwd(__nv_bfloat16* grad_src_feat,
                           __nv_bfloat16* grad_dst_feat,
                           __nv_bfloat16* grad_edge_feat,
                           __nv_bfloat16* grad_weights,
                           __nv_bfloat16* grad_wspace_feat,
                           __nv_bfloat16* grad_wspace_weight,
                           float* grad_sm_scores,
                           const __nv_bfloat16* grad_out,
                           const __nv_bfloat16* src_feat,
                           const __nv_bfloat16* dst_feat,
                           const __nv_bfloat16* edge_feat,
                           const __nv_bfloat16* attn_weights,
                           const float* sm_scores,
                           const bipartite_csc_s64_t& graph,
                           const mha_params& params,
                           int dim_node,
                           int dim_edge,
                           const cuda::stream& stream);

void mha_gat_n2n_efeat_bwd(__half* grad_src_feat,
                           __half* grad_dst_feat,
                           __half* grad_edge_feat,
                           __half* grad_weights,
                           __half* grad_wspace_feat,
                           __half* grad_wspace_weight,
                           float* grad_sm_scores,
                           const __half* grad_out,
                           const __half* src_feat,
                           const __half* dst_feat,
                           const __half* edge_feat,
                           const __half* attn_weights,
                           const float* sm_scores,
                           const bipartite_csc_s32_t& graph,
                           const mha_params& params,
                           int dim_node,
                           int dim_edge,
                           const cuda::stream& stream);
void mha_gat_n2n_efeat_bwd(__half* grad_src_feat,
                           __half* grad_dst_feat,
                           __half* grad_edge_feat,
                           __half* grad_weights,
                           __half* grad_wspace_feat,
                           __half* grad_wspace_weight,
                           float* grad_sm_scores,
                           const __half* grad_out,
                           const __half* src_feat,
                           const __half* dst_feat,
                           const __half* edge_feat,
                           const __half* attn_weights,
                           const float* sm_scores,
                           const bipartite_csc_s64_t& graph,
                           const mha_params& params,
                           int dim_node,
                           int dim_edge,
                           const cuda::stream& stream);

void mha_gat_n2n_efeat_bwd(float* grad_src_feat,
                           float* grad_dst_feat,
                           float* grad_edge_feat,
                           float* grad_weights,
                           float* grad_wspace_feat,
                           float* grad_wspace_weight,
                           float* grad_sm_scores,
                           const __half* grad_out,
                           const __half* src_feat,
                           const __half* dst_feat,
                           const __half* edge_feat,
                           const __half* attn_weights,
                           const float* sm_scores,
                           const bipartite_csc_s32_t& graph,
                           const mha_params& params,
                           int dim_node,
                           int dim_edge,
                           const cuda::stream& stream);
void mha_gat_n2n_efeat_bwd(float* grad_src_feat,
                           float* grad_dst_feat,
                           float* grad_edge_feat,
                           float* grad_weights,
                           float* grad_wspace_feat,
                           float* grad_wspace_weight,
                           float* grad_sm_scores,
                           const __half* grad_out,
                           const __half* src_feat,
                           const __half* dst_feat,
                           const __half* edge_feat,
                           const __half* attn_weights,
                           const float* sm_scores,
                           const bipartite_csc_s64_t& graph,
                           const mha_params& params,
                           int dim_node,
                           int dim_edge,
                           const cuda::stream& stream);

void mha_gat_n2n_efeat_bwd(float* grad_src_feat,
                           float* grad_dst_feat,
                           float* grad_edge_feat,
                           float* grad_weights,
                           float* grad_wspace_feat,
                           float* grad_wspace_weight,
                           float* grad_sm_scores,
                           const __nv_bfloat16* grad_out,
                           const __nv_bfloat16* src_feat,
                           const __nv_bfloat16* dst_feat,
                           const __nv_bfloat16* edge_feat,
                           const __nv_bfloat16* attn_weights,
                           const float* sm_scores,
                           const bipartite_csc_s32_t& graph,
                           const mha_params& params,
                           int dim_node,
                           int dim_edge,
                           const cuda::stream& stream);
void mha_gat_n2n_efeat_bwd(float* grad_src_feat,
                           float* grad_dst_feat,
                           float* grad_edge_feat,
                           float* grad_weights,
                           float* grad_wspace_feat,
                           float* grad_wspace_weight,
                           float* grad_sm_scores,
                           const __nv_bfloat16* grad_out,
                           const __nv_bfloat16* src_feat,
                           const __nv_bfloat16* dst_feat,
                           const __nv_bfloat16* edge_feat,
                           const __nv_bfloat16* attn_weights,
                           const float* sm_scores,
                           const bipartite_csc_s64_t& graph,
                           const mha_params& params,
                           int dim_node,
                           int dim_edge,
                           const cuda::stream& stream);

void mha_gat_n2n_efeat_bwd(__half* grad_src_feat,
                           __half* grad_dst_feat,
                           __half* grad_edge_feat,
                           float* grad_weights,
                           __half* grad_wspace_feat,
                           float* grad_wspace_weight,
                           float* grad_sm_scores,
                           const __half* grad_out,
                           const __half* src_feat,
                           const __half* dst_feat,
                           const __half* edge_feat,
                           const __half* attn_weights,
                           const float* sm_scores,
                           const bipartite_csc_s32_t& graph,
                           const mha_params& params,
                           int dim_node,
                           int dim_edge,
                           const cuda::stream& stream);
void mha_gat_n2n_efeat_bwd(__half* grad_src_feat,
                           __half* grad_dst_feat,
                           __half* grad_edge_feat,
                           float* grad_weights,
                           __half* grad_wspace_feat,
                           float* grad_wspace_weight,
                           float* grad_sm_scores,
                           const __half* grad_out,
                           const __half* src_feat,
                           const __half* dst_feat,
                           const __half* edge_feat,
                           const __half* attn_weights,
                           const float* sm_scores,
                           const bipartite_csc_s64_t& graph,
                           const mha_params& params,
                           int dim_node,
                           int dim_edge,
                           const cuda::stream& stream);

void mha_gat_n2n_efeat_bwd(__nv_bfloat16* grad_src_feat,
                           __nv_bfloat16* grad_dst_feat,
                           __nv_bfloat16* grad_edge_feat,
                           float* grad_weights,
                           __nv_bfloat16* grad_wspace_feat,
                           float* grad_wspace_weight,
                           float* grad_sm_scores,
                           const __nv_bfloat16* grad_out,
                           const __nv_bfloat16* src_feat,
                           const __nv_bfloat16* dst_feat,
                           const __nv_bfloat16* edge_feat,
                           const __nv_bfloat16* attn_weights,
                           const float* sm_scores,
                           const bipartite_csc_s32_t& graph,
                           const mha_params& params,
                           int dim_node,
                           int dim_edge,
                           const cuda::stream& stream);
void mha_gat_n2n_efeat_bwd(__nv_bfloat16* grad_src_feat,
                           __nv_bfloat16* grad_dst_feat,
                           __nv_bfloat16* grad_edge_feat,
                           float* grad_weights,
                           __nv_bfloat16* grad_wspace_feat,
                           float* grad_wspace_weight,
                           float* grad_sm_scores,
                           const __nv_bfloat16* grad_out,
                           const __nv_bfloat16* src_feat,
                           const __nv_bfloat16* dst_feat,
                           const __nv_bfloat16* edge_feat,
                           const __nv_bfloat16* attn_weights,
                           const float* sm_scores,
                           const bipartite_csc_s64_t& graph,
                           const mha_params& params,
                           int dim_node,
                           int dim_edge,
                           const cuda::stream& stream);
/** @} */

}  // namespace cugraph::ops
