def _segment_mm_backward_B(A, dC, dB, seglen):
    """Invoke the C API of the backward of segment_mm on B."""
    _CAPI_DGLKernelSEGMENTMMBackwardB(
        to_dgl_nd(A), to_dgl_nd(dC), to_dgl_nd_for_write(dB), to_dgl_nd(seglen)
    )
    return dB
Invoke the C API of the backward of segment_mm on B.
python/dgl/_sparse_ops.py