def gspmm(g, op, reduce_op, lhs_data, rhs_data):
    r"""Generalized Sparse Matrix Multiplication interface.
    It fuses two steps into one kernel.

    1. Computes messages by :attr:`op` source node and edge features.
    2. Aggregate the messages by :attr:`reduce_op` as the features on destination nodes.

    .. math::
        x_v = \psi_{(u, v, e)\in \mathcal{G}}(\rho(x_u, x_e))

    where :math:`x_v` is the returned feature on destination nodes, and :math:`x_u`,
    :math:`x_e` refers to :attr:`u`, :attr:`e` respectively. :math:`\rho` means binary
    operator :attr:`op` and :math:`\psi` means reduce operator :attr:`reduce_op`,
    :math:`\mathcal{G}` is the graph we apply gspmm on: :attr:`g`.

    Note that this function does not handle gradients.

    Parameters
    ----------
    g : DGLGraph
        The input graph.
    op : str
        The binary op's name, could be ``add``, ``sub``, ``mul``, ``div``,
        ``copy_lhs``, ``copy_rhs``.
    reduce_op : str
        Reduce operator, could be ``sum``, ``max``, ``min``, ``mean``.
    lhs_data : tensor or None
        The left operand, could be None if it's not required by the op.
    rhs_data : tensor or None
        The right operand, could be None if it's not required by the op.

    Returns
    -------
    tensor
        The result tensor.
    """
    if g._graph.number_of_etypes() == 1:
        if op not in ["copy_lhs", "copy_rhs"]:
            lhs_data, rhs_data = reshape_lhs_rhs(lhs_data, rhs_data)
        # With max and min reducers infinity will be returned for zero degree nodes
        ret = gspmm_internal(
            g._graph,
            op,
            "sum" if reduce_op == "mean" else reduce_op,
            lhs_data,
            rhs_data,
        )
    else:
        # lhs_data or rhs_data is None only in unary functions like ``copy-u`` or ``copy_e``
        lhs_data = (
            [None] * g._graph.number_of_ntypes()
            if lhs_data is None
            else lhs_data
        )
        rhs_data = (
            [None] * g._graph.number_of_etypes()
            if rhs_data is None
            else rhs_data
        )
        # TODO (Israt): Call reshape func
        lhs_and_rhs_tuple = tuple(list(lhs_data) + list(rhs_data))
        ret = gspmm_internal_hetero(
            g._graph,
            op,
            "sum" if reduce_op == "mean" else reduce_op,
            len(lhs_data),
            *lhs_and_rhs_tuple
        )
    # TODO (Israt): Add support for 'mean' in heterograph
    # divide in degrees for mean reducer.
    if reduce_op == "mean":
        ret_shape = F.shape(ret)
        deg = g.in_degrees()
        deg = F.astype(F.clamp(deg, 1, max(g.num_edges(), 1)), F.dtype(ret))
        deg_shape = (ret_shape[0],) + (1,) * (len(ret_shape) - 1)
        return ret / F.reshape(deg, deg_shape)
    else:
        return ret
Generalized Sparse Matrix Multiplication interface.
    It fuses two steps into one kernel.

    1. Computes messages by :attr:`op` source node and edge features.
    2. Aggregate the messages by :attr:`reduce_op` as the features on destination nodes.

    .. math::
        x_v = \psi_{(u, v, e)\in \mathcal{G}}(\rho(x_u, x_e))

    where :math:`x_v` is the returned feature on destination nodes, and :math:`x_u`,
    :math:`x_e` refers to :attr:`u`, :attr:`e` respectively. :math:`\rho` means binary
    operator :attr:`op` and :math:`\psi` means reduce operator :attr:`reduce_op`,
    :math:`\mathcal{G}` is the graph we apply gspmm on: :attr:`g`.

    Note that this function does not handle gradients.

    Parameters
    ----------
    g : DGLGraph
        The input graph.
    op : str
        The binary op's name, could be ``add``, ``sub``, ``mul``, ``div``,
        ``copy_lhs``, ``copy_rhs``.
    reduce_op : str
        Reduce operator, could be ``sum``, ``max``, ``min``, ``mean``.
    lhs_data : tensor or None
        The left operand, could be None if it's not required by the op.
    rhs_data : tensor or None
        The right operand, could be None if it's not required by the op.

    Returns
    -------
    tensor
        The result tensor.
    
python/dgl/ops/spmm.py