def _gsddmm_hetero(
    gidx, op, lhs_len, lhs_target="u", rhs_target="v", lhs_and_rhs_tuple=None
):
    r"""Generalized Sampled-Dense-Dense Matrix Multiplication interface."""
    lhs_tuple, rhs_tuple = (
        lhs_and_rhs_tuple[:lhs_len],
        lhs_and_rhs_tuple[lhs_len:],
    )

    use_lhs = op != "copy_rhs"
    use_rhs = op != "copy_lhs"

    # TODO (Israt): Add check - F.dtype(u) != F.dtype(e):
    # deal with scalar features.
    expand_lhs, expand_rhs = False, False
    num_ntype = gidx.number_of_ntypes()
    num_etype = gidx.number_of_etypes()
    lhs_list = (
        [None] * num_ntype if lhs_target in ["u", "v"] else [None] * num_etype
    )
    rhs_list = (
        [None] * num_ntype if rhs_target in ["u", "v"] else [None] * num_etype
    )
    out_list = [None] * gidx.number_of_etypes()

    lhs_target = target_mapping[lhs_target]
    rhs_target = target_mapping[rhs_target]

    for etid in range(gidx.number_of_etypes()):
        lhs_id = get_typeid_by_target(gidx, etid, lhs_target)
        rhs_id = get_typeid_by_target(gidx, etid, rhs_target)
        lhs = lhs_tuple[lhs_id]
        rhs = rhs_tuple[rhs_id]
        if use_lhs:
            if lhs is not None and F.ndim(lhs) == 1:
                lhs = F.unsqueeze(lhs, -1)
                expand_lhs = True
        if use_rhs:
            if rhs is not None and F.ndim(rhs) == 1:
                rhs = F.unsqueeze(rhs, -1)
                expand_rhs = True
        ctx = F.context(lhs) if use_lhs else F.context(rhs)
        dtype = F.dtype(lhs) if use_lhs else F.dtype(rhs)
        lhs_shp = F.shape(lhs) if use_lhs else (0,)
        rhs_shp = F.shape(rhs) if use_rhs else (0,)
        lhs_list[lhs_id] = lhs if use_lhs else None
        rhs_list[rhs_id] = rhs if use_rhs else None
        out_shp = (gidx.num_edges(etid),) + infer_broadcast_shape(
            op, lhs_shp[1:], rhs_shp[1:]
        )
        out_list[etid] = F.empty(out_shp, dtype, ctx)
    if gidx.num_edges(0) > 0:
        _CAPI_DGLKernelSDDMMHetero(
            gidx,
            op,
            [to_dgl_nd(lhs) for lhs in lhs_list],
            [to_dgl_nd(rhs) for rhs in rhs_list],
            [to_dgl_nd_for_write(out) for out in out_list],
            lhs_target,
            rhs_target,
        )

    for l in range(gidx.number_of_etypes()):
        # Replace None by empty tensor. Forward func doesn't accept None in tuple.
        e = out_list[l]
        e = F.tensor([]) if e is None else e
        if (expand_lhs or not use_lhs) and (expand_rhs or not use_rhs):
            e = F.squeeze(e, -1)
        out_list[l] = e
    out = tuple(out_list)
    return out
Generalized Sampled-Dense-Dense Matrix Multiplication interface.
python/dgl/_sparse_ops.py