def _gspmm_hetero(gidx, op, reduce_op, u_len, u_and_e_tuple):
    r"""Generalized Sparse Matrix Multiplication interface on heterogeneous graphs.
    It handles multiple node and edge types of the graph. For each edge type, it takes
    the result of :attr:`op` on source node feature and edge feature, and leads to a
    message on edge. Then it aggregates the message by :attr:`reduce_op` on the destination
    nodes of the etype.

    .. 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
    ----------
    gidx : HeteroGraphIndex
        The input graph index.
    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``.
    u_len : int
        The number of tensors in ``u`` (source node features)
    u_and_e_tuple : Tuple of tensors
        Tuple of source nodes' features and edges' features. ``u_and_e_tuple[:u_len]``
        stores the source nodes's features of all source node types. ``u_and_e_tuple[u_len:]``
        stores the edges's features of all the edge types.
        The source nodes' features of the soruce node types could be None if op is ``copy_rhs``.
        The edges' features of the edge types could be None if op is ``copy_lhs``.

    Returns
    -------
    tuple
        The returned tuple is composed of two elements:
        - The first element refers to the tuple of result tensors.
        - The second element refers to a tuple composed of arg_u and arg_e
          (which is useful when reducer is `min`/`max`).

    Notes
    -----
    This function does not handle gradients.
    """
    u_tuple, e_tuple = u_and_e_tuple[:u_len], u_and_e_tuple[u_len:]
    use_u = op != "copy_rhs"
    use_e = op != "copy_lhs"
    # TODO (Israt): Add check - F.dtype(u) != F.dtype(e):

    # deal with scalar features.
    expand_u, expand_e = False, False
    num_ntypes = gidx.number_of_ntypes()
    num_etypes = gidx.number_of_etypes()
    list_u = [None] * num_ntypes
    list_v = [None] * num_ntypes
    list_e = [None] * num_etypes
    list_arg_u_nd = [None] * num_ntypes
    list_arg_u = [None] * num_ntypes
    list_arg_u_ntype_nd = [None] * num_ntypes
    list_arg_u_ntype = [None] * num_ntypes
    # TODO(Israt): double check ntype or etype
    list_arg_e_nd = [None] * num_ntypes
    list_arg_e = [None] * num_ntypes
    list_arg_e_etype_nd = [None] * num_ntypes
    list_arg_e_etype = [None] * num_ntypes

    use_cmp = reduce_op in ["max", "min"]
    idtype = getattr(F, gidx.dtype)

    for etid in range(num_etypes):
        src_id, dst_id = gidx.metagraph.find_edge(etid)
        u = u_tuple[src_id] if use_u else None
        e = e_tuple[etid] if use_e else None
        if use_u:
            if u is not None and F.ndim(u) == 1:
                u = F.unsqueeze(u, -1)
                expand_u = True
            list_u[src_id] = u if use_u else None
        if use_e:
            if e is not None and F.ndim(e) == 1:
                e = F.unsqueeze(e, -1)
                expand_e = True
            list_e[etid] = e if use_e else None
        ctx = (
            F.context(u) if use_u else F.context(e)
        )  # TODO(Israt): Put outside of loop
        dtype = (
            F.dtype(u) if use_u else F.dtype(e)
        )  # TODO(Israt): Put outside of loop
        u_shp = F.shape(u) if use_u else (0,)
        e_shp = F.shape(e) if use_e else (0,)
        v_shp = (gidx.num_nodes(dst_id),) + infer_broadcast_shape(
            op, u_shp[1:], e_shp[1:]
        )
        list_v[dst_id] = F.zeros(v_shp, dtype, ctx)
        if use_cmp:
            if use_u:
                list_arg_u[dst_id] = F.zeros(v_shp, idtype, ctx)
                list_arg_u_ntype[dst_id] = F.zeros(v_shp, idtype, ctx)
            if use_e:
                list_arg_e[dst_id] = F.zeros(v_shp, idtype, ctx)
                list_arg_e_etype[dst_id] = F.zeros(v_shp, idtype, ctx)
        list_arg_u_nd[dst_id] = to_dgl_nd_for_write(list_arg_u[dst_id])
        list_arg_u_ntype_nd[dst_id] = to_dgl_nd_for_write(
            list_arg_u_ntype[dst_id]
        )
        list_arg_e_nd[dst_id] = to_dgl_nd_for_write(list_arg_e[dst_id])
        list_arg_e_etype_nd[dst_id] = to_dgl_nd_for_write(
            list_arg_e_etype[dst_id]
        )

    if gidx.num_edges(0) > 0:
        _CAPI_DGLKernelSpMMHetero(
            gidx,
            op,
            reduce_op,
            [to_dgl_nd(u_i) for u_i in list_u],
            [to_dgl_nd(e_i) for e_i in list_e],
            [to_dgl_nd_for_write(v_i) for v_i in list_v],
            list_arg_u_nd,
            list_arg_e_nd,
            list_arg_u_ntype_nd,
            list_arg_e_etype_nd,
        )
    for l, arg_u_nd in enumerate(list_arg_u_nd):
        # TODO(Israt): l or src_id as index of lhs
        list_arg_u[l] = (
            None
            if list_arg_u[l] is None
            else F.zerocopy_from_dgl_ndarray(arg_u_nd)
        )
        if list_arg_u[l] is not None and expand_u and use_cmp:
            list_arg_u[l] = F.squeeze(list_arg_u[l], -1)
    for l, arg_e_nd in enumerate(list_arg_e_nd):
        list_arg_e[l] = (
            None
            if list_arg_e[l] is None
            else F.zerocopy_from_dgl_ndarray(arg_e_nd)
        )
        if list_arg_e[l] is not None and expand_e and use_cmp:
            list_arg_e[l] = F.squeeze(list_arg_e[l], -1)
    for l, arg_u_ntype_nd in enumerate(list_arg_u_ntype_nd):
        list_arg_u_ntype[l] = (
            None
            if arg_u_ntype_nd is None
            else F.zerocopy_from_dgl_ndarray(arg_u_ntype_nd)
        )
    for l, arg_e_etype_nd in enumerate(list_arg_e_etype_nd):
        list_arg_e_etype[l] = (
            None
            if arg_e_etype_nd is None
            else F.zerocopy_from_dgl_ndarray(arg_e_etype_nd)
        )
    # To deal with scalar node/edge features.
    for l in range(num_ntypes):
        # replace None by empty tensor. Forward func doesn't accept None in tuple.
        v = list_v[l]
        v = F.tensor([]) if v is None else v
        if (expand_u or not use_u) and (expand_e or not use_e):
            v = F.squeeze(v, -1)  # To deal with scalar node/edge features.
        list_v[l] = v
    out = tuple(list_v)
    return out, (list_arg_u, list_arg_e, list_arg_u_ntype, list_arg_e_etype)
Generalized Sparse Matrix Multiplication interface on heterogeneous graphs.
    It handles multiple node and edge types of the graph. For each edge type, it takes
    the result of :attr:`op` on source node feature and edge feature, and leads to a
    message on edge. Then it aggregates the message by :attr:`reduce_op` on the destination
    nodes of the etype.

    .. 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
    ----------
    gidx : HeteroGraphIndex
        The input graph index.
    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``.
    u_len : int
        The number of tensors in ``u`` (source node features)
    u_and_e_tuple : Tuple of tensors
        Tuple of source nodes' features and edges' features. ``u_and_e_tuple[:u_len]``
        stores the source nodes's features of all source node types. ``u_and_e_tuple[u_len:]``
        stores the edges's features of all the edge types.
        The source nodes' features of the soruce node types could be None if op is ``copy_rhs``.
        The edges' features of the edge types could be None if op is ``copy_lhs``.

    Returns
    -------
    tuple
        The returned tuple is composed of two elements:
        - The first element refers to the tuple of result tensors.
        - The second element refers to a tuple composed of arg_u and arg_e
          (which is useful when reducer is `min`/`max`).

    Notes
    -----
    This function does not handle gradients.
    
python/dgl/_sparse_ops.py