Method forward

python/dgl/backend/pytorch/sparse.py:687–718

Forward function. Pseudo-code: .. code:: python score = dgl.EData(g, score) score_max = score.dst_max() # of type dgl.NData score = score - score_max # edge_sub_dst, ret dgl.EData score_sum = score.dst_sum() # of type dgl.NData

(ctx, gidx, score, eids, norm_by)

class EdgeSoftmax(th.autograd.Function):
    @staticmethod
    def forward(ctx, gidx, score, eids, norm_by):
        """Forward function.

        Pseudo-code:

        .. code:: python

            score = dgl.EData(g, score)
            score_max = score.dst_max()  # of type dgl.NData
            score = score - score_max  # edge_sub_dst, ret dgl.EData
            score_sum = score.dst_sum()  # of type dgl.NData
            out = score / score_sum    # edge_div_dst, ret dgl.EData
            return out.data
        """
        # remember to save the graph to backward cache before making it
        # a local variable
        if not is_all(eids):
            gidx = gidx.edge_subgraph([eids], True).graph
        if norm_by == "src":
            gidx = gidx.reverse()
        # Note: Now _edge_softmax_forward op only supports CPU
        # TODO(Zhejiang): We will support GPU in the future
        if score.is_cuda:
            score_max = _gspmm(gidx, "copy_rhs", "max", None, score)[0]
            score = th.exp(_gsddmm(gidx, "sub", score, score_max, "e", "v"))
            score_sum = _gspmm(gidx, "copy_rhs", "sum", None, score)[0]
            out = _gsddmm(gidx, "div", score, score_sum, "e", "v")
        else:
            out = _edge_softmax_forward(gidx, score, "copy_rhs")
        ctx.backward_cache = gidx
        ctx.save_for_backward(out)
        return out

    @staticmethod
    def backward(ctx, grad_out):