# Copyright (c) 2023-2024, NVIDIA CORPORATION.
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import torch
from typing import Optional, Tuple, Union
from torch.autograd.function import once_differentiable

from ..utils import (
    maybe_empty_like,
    maybe_empty,
    get_tensor_meta_data,
    maybe_zeros,
    maybe_zeros_like,
)
from ..graph import CSC
from pylibcugraphops.operators import (
    AggOp,
    agg_concat_n2n_fwd,
    agg_concat_n2n_bwd,
    agg_concat_e2n_fwd,
    agg_concat_e2n_bwd,
    agg_concat_n2n_e2n_fwd,
    agg_concat_n2n_e2n_bwd,
    agg_concat_weighted_n2n_fwd,
    agg_concat_weighted_n2n_bwd,
    agg_concat_n2n_bipartite_fwd,
    agg_concat_n2n_bipartite_bwd,
    agg_concat_e2n_bipartite_fwd,
    agg_concat_e2n_bipartite_bwd,
    agg_concat_n2n_e2n_bipartite_fwd,
    agg_concat_n2n_e2n_bipartite_bwd,
    agg_concat_weighted_n2n_bipartite_fwd,
    agg_concat_weighted_n2n_bipartite_bwd,
)

_aggr_ops = ["max", "mean", "min", "sum"]


def agg_concat_n2n(
    feat: Union[torch.Tensor, Tuple[torch.Tensor]],
    graph: CSC,
    aggr: str = "sum",
    edge_weight: Optional[torch.Tensor] = None,
) -> torch.Tensor:
    """
    PyTorch autograd function for simple aggregation using node features in
    an node-to-node reduction (n2n) while concatenating the original features of
    output at the end (agg_concat).

    Parameters
    ----------
    feat : torch.Tensor | Tuple[torch.Tensor]
        The input node features.
        Shape: ``(n_src_nodes, dim_in)`` if graph is non-bipartite.
        Shape: ``(n_src_nodes, dim_in)`` and ``(n_dst_nodes, dim_concat)``
        if graph is bipartite.

    graph : CSC
        The graph used for the operation.

    aggr : str, default="sum"
        Aggregation operation. Choose from ``"max"``, ``"mean"``, ``"min"``, ``"sum"``.

    edge_weight : Optional[torch.Tensor], default=None
        When passing additional edge weights, the aggregation is done in
        a weighted fashion. Note that for min/max the result after
        weighting features with those edge weights is returned.
        Edge weights are indexed in the same fashion as the indices in ``graph``.
        Shape: ``(n_edges, )``.

    Returns
    -------
    output : torch.Tensor
        The aggregation output.
        Shape: ``(n_dst_nodes, 2 * dim_in)``.
    """

    if aggr not in _aggr_ops:
        raise ValueError(
            f"Supported aggregation operations: {_aggr_ops}, " f"but got '{aggr}'."
        )
    aggr = getattr(AggOp, aggr.capitalize())

    if graph.is_bipartite:
        assert isinstance(feat, Tuple) and len(feat) == 2, (
            f"Expected tuple of source and destination node features,"
            f" got {type(feat)} instead."
        )
        return _agg_concat_n2n_bipartite_autograd.apply(
            feat[0], feat[1], edge_weight, graph, aggr
        )

    assert isinstance(
        feat, torch.Tensor
    ), f"Expected torch.Tensor of node features, got {type(feat)} instead."
    return _agg_concat_n2n_autograd.apply(feat, edge_weight, graph, aggr)


def agg_concat_e2n(
    node_feat: torch.Tensor,
    edge_feat: torch.Tensor,
    graph: CSC,
    aggr: str = "sum",
) -> torch.Tensor:
    r"""PyTorch autograd function for simple aggregation using edge features in
    while concatenating the original features of output at the end (agg_concat).

    Parameters
    ----------
    node_feat : torch.Tensor
        The input node features.
        Shape: ``(n_dst_nodes, dim_node)``.

    edge_feat : torch.Tensor
        The input edge features.
        Shape: ``(n_edges, dim_edge)``.

    graph : CSC
        The graph used for the operation.

    aggr : str, default="sum"
        Aggregation operation. Choose from ``"max"``, ``"mean"``, ``"min"``, ``"sum"``.

    Returns
    -------
    output : torch.Tensor
        The aggregation output.
        Shape: ``(n_dst_nodes, dim_node + dim_edge)``.
    """

    if aggr not in _aggr_ops:
        raise ValueError(
            f"Supported aggregation operations: {_aggr_ops}, " f"but got '{aggr}'."
        )
    aggr = getattr(AggOp, aggr.capitalize())

    return _agg_concat_e2n_autograd.apply(edge_feat, node_feat, graph, aggr)


def agg_concat_n2n_e2n(
    feat: Union[torch.Tensor, Tuple[torch.Tensor]],
    edge_feat: torch.Tensor,
    graph: CSC,
    aggr: str = "sum",
) -> torch.Tensor:
    r"""PyTorch autograd function for simple aggregation using node features in
    an node-to-node reduction (n2n) and edge features in and edge-to-node
    reduction (e2n) while concatenating the original features of output at the
    end (agg_concat).

    Parameters
    ----------
    node_feat : torch.Tensor
        The input node features.
        Shape: ``(n_src_nodes, dim_in)`` if graph is non-bipartite.
        Shape: ``(n_src_nodes, dim_in)`` and ``(n_dst_nodes, dim_concat)``
        if graph is bipartite.

    edge_feat : torch.Tensor
        The input edge features.
        Shape: ``(n_edges, dim_edge)``.

    graph : CSC
        The graph used for the operation.

    aggr : str, default="sum"
        Aggregation operation. Choose from ``"max"``, ``"mean"``, ``"min"``, ``"sum"``.

    Returns
    -------
    output : torch.Tensor
        The aggregation output.
        Shape: ``(n_dst_nodes, dim_node + dim_edge + dim_node)``.
    """

    if aggr not in _aggr_ops:
        raise ValueError(
            f"Supported aggregation operations: {_aggr_ops}, " f"but got '{aggr}'."
        )
    aggr = getattr(AggOp, aggr.capitalize())

    if graph.is_bipartite:
        assert isinstance(feat, Tuple) and len(feat) == 2, (
            f"Expected tuple of source and destination node features,"
            f" got {type(feat)} instead."
        )
        return _agg_concat_n2n_e2n_bipartite_autograd.apply(
            feat[0],
            edge_feat,
            feat[1],
            graph,
            aggr,
        )

    assert isinstance(
        feat, torch.Tensor
    ), f"Expected torch.Tensor of node features, got {type(feat)} instead."
    return _agg_concat_n2n_e2n_autograd.apply(feat, edge_feat, graph, aggr)


class _agg_concat_n2n_autograd(torch.autograd.Function):
    """
    Custom autograd function for simple aggregation using node features in
    an node-to-node reduction (n2n) on a graph while concatenating the
    original features of output at the end (agg_concat).
    """

    @staticmethod
    def forward(ctx, feat, edge_weight, graph, aggr):
        dim_in = feat.size(-1)
        agg_output = torch.empty(
            graph.num_dst_nodes, 2 * dim_in, dtype=feat.dtype, device=feat.device
        )

        node_deg = None
        out_pos = None
        if graph.num_dst_nodes > 0:
            if aggr in (AggOp.Max, AggOp.Min):
                idx_type = graph.graph_csc.__class__.__name__.split("_")[-1]
                out_pos = torch.empty(
                    graph.num_dst_nodes,
                    dim_in,
                    dtype=getattr(torch, idx_type),
                    device=feat.device,
                )

            stream = torch.cuda.current_stream().cuda_stream
            feat = feat.detach().contiguous()

            if edge_weight is None:
                agg_concat_n2n_fwd(
                    agg_output, feat, graph.graph_csc, aggr, out_pos, stream_id=stream
                )

            else:
                edge_weight = edge_weight.detach().contiguous()
                node_deg = torch.empty(
                    graph.num_dst_nodes,
                    device=edge_weight.device,
                    dtype=edge_weight.dtype,
                )

                agg_concat_weighted_n2n_fwd(
                    agg_output,
                    feat,
                    edge_weight,
                    graph.graph_csc,
                    aggr,
                    out_pos,
                    node_deg,
                    stream_id=stream,
                )

        ctx.graph = graph
        ctx.aggr = aggr
        ctx.dim_node = dim_in
        ctx.feat_meta_data = get_tensor_meta_data(feat)

        saved_agg_output = (
            agg_output.detach().contiguous()
            if (edge_weight is not None) and (aggr == AggOp.Mean)
            else None
        )
        saved_feat = feat if edge_weight is not None else None
        ctx.save_for_backward(
            saved_feat, edge_weight, node_deg, out_pos, saved_agg_output
        )

        return agg_output

    @staticmethod
    @once_differentiable
    def backward(ctx, grad_output):
        feat, edge_weight, node_deg, out_pos, agg_output = ctx.saved_tensors
        needs_grad_feat, needs_grad_weight, _, _ = ctx.needs_input_grad
        if not (needs_grad_feat or needs_grad_weight):
            return None, None, None, None

        grad_output = grad_output.detach().contiguous()

        if ctx.graph.num_dst_nodes > 0:
            grad_feat = maybe_empty(
                **ctx.feat_meta_data,
                create_tensor=needs_grad_feat,
            )
            # captures both having edge_weight=None and
            # needs_grad_weight=False with edge_weight!=None
            grad_edge_weight = maybe_empty_like(edge_weight, needs_grad_weight)
            stream = torch.cuda.current_stream().cuda_stream

            if edge_weight is None:
                agg_concat_n2n_bwd(
                    grad_feat,
                    grad_output,
                    ctx.graph.graph_csc,
                    ctx.dim_node,
                    ctx.aggr,
                    out_pos,
                    stream_id=stream,
                )

            else:
                agg_concat_weighted_n2n_bwd(
                    grad_feat,
                    grad_edge_weight,
                    grad_output,
                    edge_weight,
                    feat,
                    ctx.graph.graph_csc,
                    ctx.dim_node,
                    ctx.aggr,
                    agg_output,
                    out_pos,
                    node_deg,
                    stream_id=stream,
                )
        else:
            grad_feat = maybe_zeros(
                **ctx.feat_meta_data,
                create_tensor=needs_grad_feat,
            )
            grad_edge_weight = maybe_zeros_like(edge_weight, needs_grad_weight)

        return grad_feat, grad_edge_weight, None, None


class _agg_concat_n2n_bipartite_autograd(torch.autograd.Function):
    """
    Custom autograd function for simple aggregation using node features in
    an node-to-node reduction (n2n) on a graph while concatenating the
    original features of output at the end (agg_concat).
    """

    @staticmethod
    def forward(ctx, src_feat, concat_feat, edge_weight, graph, aggr):
        dim_in = src_feat.size(-1)
        dim_concat = concat_feat.size(-1)

        agg_output = torch.empty(
            graph.num_dst_nodes,
            dim_in + dim_concat,
            dtype=src_feat.dtype,
            device=src_feat.device,
        )
        node_deg = None
        out_pos = None
        if graph.num_dst_nodes > 0:
            if aggr in (AggOp.Max, AggOp.Min):
                idx_type = graph.graph_csc.__class__.__name__.split("_")[-1]
                out_pos = torch.empty(
                    graph.num_dst_nodes,
                    dim_in,
                    dtype=getattr(torch, idx_type),
                    device=src_feat.device,
                )

            stream = torch.cuda.current_stream().cuda_stream

            src_feat = src_feat.detach().contiguous()
            concat_feat = concat_feat.detach().contiguous()
            if edge_weight is None:
                agg_concat_n2n_bipartite_fwd(
                    agg_output,
                    src_feat,
                    concat_feat,
                    graph.graph_csc,
                    aggr,
                    out_pos,
                    stream_id=stream,
                )

            else:
                edge_weight = edge_weight.detach().contiguous()
                node_deg = torch.empty(
                    graph.num_dst_nodes,
                    device=edge_weight.device,
                    dtype=edge_weight.dtype,
                )

                agg_concat_weighted_n2n_bipartite_fwd(
                    agg_output,
                    src_feat,
                    concat_feat,
                    edge_weight,
                    graph.graph_csc,
                    aggr,
                    out_pos,
                    node_deg,
                    stream_id=stream,
                )

        ctx.graph = graph
        ctx.aggr = aggr
        ctx.dim_node = dim_in
        ctx.dim_concat = dim_concat
        ctx.feat_meta_data = get_tensor_meta_data(src_feat)
        ctx.concat_meta_data = get_tensor_meta_data(concat_feat)

        saved_agg_output = (
            agg_output.detach().contiguous()
            if (edge_weight is not None) and (aggr == AggOp.Mean)
            else None
        )
        saved_feat = src_feat if edge_weight is not None else None
        ctx.save_for_backward(
            saved_feat, edge_weight, node_deg, out_pos, saved_agg_output
        )

        return agg_output

    @staticmethod
    @once_differentiable
    def backward(ctx, grad_output):
        feat, edge_weight, node_deg, out_pos, agg_output = ctx.saved_tensors
        (
            needs_grad_feat,
            needs_grad_concat,
            needs_grad_weight,
            _,
            _,
        ) = ctx.needs_input_grad
        if not (needs_grad_feat or needs_grad_weight or needs_grad_concat):
            return None, None, None, None, None

        grad_output = grad_output.detach().contiguous()

        if ctx.graph.num_dst_nodes > 0:
            grad_feat = maybe_empty(
                **ctx.feat_meta_data,
                create_tensor=needs_grad_feat,
            )
            grad_concat = maybe_empty(
                **ctx.concat_meta_data,
                create_tensor=needs_grad_concat,
            )

            # captures both having edge_weight=None and
            # needs_grad_weight=False with edge_weight!=None
            grad_edge_weight = maybe_empty_like(edge_weight, needs_grad_weight)
            stream = torch.cuda.current_stream().cuda_stream

            if edge_weight is None:
                agg_concat_n2n_bipartite_bwd(
                    grad_feat,
                    grad_concat,
                    grad_output,
                    ctx.graph.graph_csc,
                    ctx.dim_node,
                    ctx.dim_concat,
                    ctx.aggr,
                    out_pos,
                    stream_id=stream,
                )

            else:
                agg_concat_weighted_n2n_bipartite_bwd(
                    grad_feat,
                    grad_concat,
                    grad_edge_weight,
                    grad_output,
                    edge_weight,
                    feat,
                    ctx.graph.graph_csc,
                    ctx.dim_node,
                    ctx.dim_concat,
                    ctx.aggr,
                    agg_output,
                    out_pos,
                    node_deg,
                    stream_id=stream,
                )
        else:
            grad_feat = maybe_zeros(
                **ctx.feat_meta_data,
                create_tensor=needs_grad_feat,
            )
            grad_concat = maybe_zeros(
                **ctx.concat_meta_data,
                create_tensor=needs_grad_concat,
            )
            grad_edge_weight = maybe_zeros_like(edge_weight, needs_grad_weight)

        return grad_feat, grad_concat, grad_edge_weight, None, None


class _agg_concat_e2n_autograd(torch.autograd.Function):
    r"""Custom autograd function for simple aggregation using edge features in
    an edge-to-node reduction (e2n) on a graph wile concatenating the
    original features of output at the end (agg_concat)."""

    @staticmethod
    def forward(ctx, edge_feat, node_feat, graph, aggr):
        dim_node = node_feat.size(-1)
        dim_edge = edge_feat.size(-1)
        agg_output = torch.empty(
            graph.num_dst_nodes,
            dim_edge + dim_node,
            dtype=edge_feat.dtype,
            device=edge_feat.device,
        )

        out_pos = None
        if graph.num_dst_nodes > 0:
            if aggr in (AggOp.Max, AggOp.Min):
                idx_type = graph.graph_csc.__class__.__name__.split("_")[-1]
                out_pos = torch.empty(
                    graph.num_dst_nodes,
                    dim_edge,
                    dtype=getattr(torch, idx_type),
                    device=edge_feat.device,
                )

            stream = torch.cuda.current_stream().cuda_stream
            node_feat = node_feat.detach().contiguous()
            edge_feat = edge_feat.detach().contiguous()
            if graph.is_bipartite:
                agg_concat_e2n_bipartite_fwd(
                    agg_output,
                    edge_feat,
                    node_feat,
                    graph.graph_csc,
                    aggr,
                    out_pos,
                    stream_id=stream,
                )
            else:
                agg_concat_e2n_fwd(
                    agg_output,
                    edge_feat,
                    node_feat,
                    graph.graph_csc,
                    aggr,
                    out_pos,
                    stream_id=stream,
                )

        ctx.graph = graph
        ctx.aggr = aggr
        ctx.dim_node = dim_node
        ctx.dim_edge = dim_edge

        ctx.edge_feat_meta_data = get_tensor_meta_data(edge_feat)
        ctx.node_feat_meta_data = get_tensor_meta_data(node_feat)

        ctx.save_for_backward(out_pos)

        return agg_output

    @staticmethod
    @once_differentiable
    def backward(ctx, grad_output):
        grad_output = grad_output.detach().contiguous()
        (out_pos,) = ctx.saved_tensors
        needs_grad_edge, needs_grad_node, _, _ = ctx.needs_input_grad
        if not (needs_grad_node or needs_grad_edge):
            return None, None, None, None

        grad_output = grad_output.detach().contiguous()

        if ctx.graph.num_dst_nodes > 0:
            grad_node_feat = maybe_empty(
                **ctx.node_feat_meta_data,
                create_tensor=needs_grad_node,
            )
            grad_edge_feat = maybe_empty(
                **ctx.edge_feat_meta_data,
                create_tensor=needs_grad_edge,
            )

            stream = torch.cuda.current_stream().cuda_stream

            if ctx.graph.is_bipartite:
                agg_concat_e2n_bipartite_bwd(
                    grad_edge_feat,
                    grad_node_feat,
                    grad_output,
                    ctx.graph.graph_csc,
                    ctx.dim_edge,
                    ctx.dim_node,
                    ctx.aggr,
                    out_pos,
                    stream_id=stream,
                )

            else:
                agg_concat_e2n_bwd(
                    grad_edge_feat,
                    grad_node_feat,
                    grad_output,
                    ctx.graph.graph_csc,
                    ctx.dim_edge,
                    ctx.dim_node,
                    ctx.aggr,
                    out_pos,
                    stream_id=stream,
                )
        else:
            grad_node_feat = maybe_zeros(
                **ctx.node_feat_meta_data,
                create_tensor=needs_grad_node,
            )
            grad_edge_feat = maybe_zeros(
                **ctx.edge_feat_meta_data,
                create_tensor=needs_grad_edge,
            )

        return grad_edge_feat, grad_node_feat, None, None


class _agg_concat_n2n_e2n_autograd(torch.autograd.Function):
    r"""Custom autograd function for simple aggregation using node features in
    an node-to-node reduction (n2n) and edge features in
    an edge-to-node reduction (e2n) on a graph while concatenating the
    original features of output at the end (agg_concat)."""

    @staticmethod
    def forward(ctx, node_feat, edge_feat, graph, aggr):
        dim_node = node_feat.size(-1)
        dim_edge = edge_feat.size(-1)
        agg_output = torch.empty(
            graph.num_dst_nodes,
            dim_node + dim_edge + dim_node,
            dtype=edge_feat.dtype,
            device=edge_feat.device,
        )

        out_pos = None
        if graph.num_dst_nodes > 0:
            if aggr in (AggOp.Max, AggOp.Min):
                idx_type = graph.graph_csc.__class__.__name__.split("_")[-1]
                out_pos = torch.empty(
                    graph.num_dst_nodes,
                    dim_node + dim_edge,
                    dtype=getattr(torch, idx_type),
                    device=edge_feat.device,
                )

            stream = torch.cuda.current_stream().cuda_stream

            node_feat = node_feat.detach().contiguous()
            edge_feat = edge_feat.detach().contiguous()
            agg_concat_n2n_e2n_fwd(
                agg_output,
                node_feat,
                edge_feat,
                graph.graph_csc,
                aggr,
                out_pos,
                stream_id=stream,
            )

        ctx.graph = graph
        ctx.dim_node = dim_node
        ctx.dim_edge = dim_edge
        ctx.aggr = aggr
        ctx.node_feat_meta_data = get_tensor_meta_data(node_feat)
        ctx.edge_feat_meta_data = get_tensor_meta_data(edge_feat)

        ctx.save_for_backward(out_pos)

        return agg_output

    @staticmethod
    @once_differentiable
    def backward(ctx, grad_output):
        grad_output = grad_output.detach().contiguous()
        (out_pos,) = ctx.saved_tensors
        needs_grad_node, needs_grad_edge, _, _ = ctx.needs_input_grad
        if not (needs_grad_node or needs_grad_edge):
            return None, None, None, None

        grad_output = grad_output.detach().contiguous()

        if ctx.graph.num_dst_nodes > 0:
            grad_node_feat = maybe_empty(
                **ctx.node_feat_meta_data, create_tensor=needs_grad_node
            )
            grad_edge_feat = maybe_empty(
                **ctx.edge_feat_meta_data, create_tensor=needs_grad_edge
            )

            stream = torch.cuda.current_stream().cuda_stream

            agg_concat_n2n_e2n_bwd(
                grad_node_feat,
                grad_edge_feat,
                grad_output,
                ctx.graph.graph_csc,
                ctx.dim_node,
                ctx.dim_edge,
                ctx.aggr,
                out_pos,
                stream_id=stream,
            )
        else:
            grad_node_feat = maybe_zeros(
                **ctx.node_feat_meta_data, create_tensor=needs_grad_node
            )
            grad_edge_feat = maybe_zeros(
                **ctx.edge_feat_meta_data, create_tensor=needs_grad_edge
            )

        return grad_node_feat, grad_edge_feat, None, None


class _agg_concat_n2n_e2n_bipartite_autograd(torch.autograd.Function):
    r"""Custom autograd function for simple aggregation using node features in
    an node-to-node reduction (n2n) and edge features in
    an edge-to-node reduction (e2n) on a graph while concatenating the
    original features of output at the end (agg_concat)."""

    @staticmethod
    def forward(ctx, src_feat, edge_feat, dst_feat, graph, aggr):
        dim_node = src_feat.size(-1)
        dim_edge = edge_feat.size(-1)
        dim_concat = dst_feat.size(-1)

        agg_output = torch.empty(
            graph.num_dst_nodes,
            dim_node + dim_edge + dim_concat,
            dtype=edge_feat.dtype,
            device=edge_feat.device,
        )

        out_pos = None
        if graph.num_dst_nodes > 0:
            if aggr in (AggOp.Max, AggOp.Min):
                idx_type = graph.graph_csc.__class__.__name__.split("_")[-1]
                out_pos = torch.empty(
                    graph.num_dst_nodes,
                    dim_node + dim_edge,
                    dtype=getattr(torch, idx_type),
                    device=edge_feat.device,
                )

            stream = torch.cuda.current_stream().cuda_stream

            src_feat = src_feat.detach().contiguous()
            edge_feat = edge_feat.detach().contiguous()
            dst_feat = dst_feat.detach().contiguous()

            agg_concat_n2n_e2n_bipartite_fwd(
                agg_output,
                src_feat,
                edge_feat,
                dst_feat,
                graph.graph_csc,
                aggr,
                out_pos,
                stream_id=stream,
            )

        ctx.graph = graph
        ctx.dim_node = dim_node
        ctx.dim_edge = dim_edge
        ctx.dim_concat = dim_concat
        ctx.aggr = aggr
        ctx.src_feat_meta_data = get_tensor_meta_data(src_feat)
        ctx.edge_feat_meta_data = get_tensor_meta_data(edge_feat)
        ctx.concat_feat_meta_data = get_tensor_meta_data(dst_feat)

        ctx.save_for_backward(out_pos)

        return agg_output

    @staticmethod
    @once_differentiable
    def backward(ctx, grad_output):
        grad_output = grad_output.detach().contiguous()
        (out_pos,) = ctx.saved_tensors
        needs_grad_node, needs_grad_edge, needs_grad_concat, _, _ = ctx.needs_input_grad
        if not (needs_grad_node or needs_grad_edge or needs_grad_concat):
            return None, None, None, None, None

        grad_output = grad_output.detach().contiguous()

        if ctx.graph.num_dst_nodes > 0:
            grad_src_feat = maybe_empty(
                **ctx.src_feat_meta_data, create_tensor=needs_grad_node
            )
            grad_edge_feat = maybe_empty(
                **ctx.edge_feat_meta_data, create_tensor=needs_grad_edge
            )
            grad_concat_feat = maybe_empty(
                **ctx.concat_feat_meta_data, create_tensor=needs_grad_concat
            )

            stream = torch.cuda.current_stream().cuda_stream

            agg_concat_n2n_e2n_bipartite_bwd(
                grad_src_feat,
                grad_edge_feat,
                grad_concat_feat,
                grad_output,
                ctx.graph.graph_csc,
                ctx.dim_node,
                ctx.dim_edge,
                ctx.dim_concat,
                ctx.aggr,
                out_pos,
                stream_id=stream,
            )
        else:
            grad_src_feat = maybe_zeros(
                **ctx.src_feat_meta_data, create_tensor=needs_grad_node
            )
            grad_edge_feat = maybe_zeros(
                **ctx.edge_feat_meta_data, create_tensor=needs_grad_edge
            )
            grad_concat_feat = maybe_zeros(
                **ctx.concat_feat_meta_data, create_tensor=needs_grad_concat
            )

        return grad_src_feat, grad_edge_feat, grad_concat_feat, None, None
