def _gather_mm(A, B, out, idx_a=None, idx_b=None):
    r"""Invoke the C API of the gather_mm operator."""
    _CAPI_DGLKernelGATHERMM(
        to_dgl_nd(A),
        to_dgl_nd(B),
        to_dgl_nd_for_write(out),
        to_dgl_nd(idx_a),
        to_dgl_nd(idx_b),
    )
    return out
Invoke the C API of the gather_mm operator.
python/dgl/_sparse_ops.py