Method forward

python/dgl/backend/pytorch/sparse.py:830–834

(ctx, op, x, offsets)

class SegmentReduce(th.autograd.Function):
    @staticmethod
    def forward(ctx, op, x, offsets):
        y, arg = _segment_reduce(op, x, offsets)
        ctx.save_for_backward(arg, offsets)
        ctx.backward_cache = op
        return y

    @staticmethod
    def backward(ctx, dy):