fc9d30fae42e3849d80b0ebf764b6e58c3198f3b,python/dgl/nn/mxnet/softmax.py,EdgeSoftmax,backward,#EdgeSoftmax#Any#,57

Before Change

        out, = self.saved_tensors  // pylint: disable=access-member-before-definition, unpacking-non-sequence
        // clear saved tensors explicitly
        self.saved_tensors = None
        out_name = utils.get_edata_name(g, "out")
        accum_name = utils.get_ndata_name(g, "accum")
        grad_score_name = utils.get_edata_name(g, "grad_score")
        g.edata[out_name] = out
        g.edata[grad_score_name] = out * grad_out
        g.update_all(fn.copy_e(grad_score_name, "m"), fn.sum("m", accum_name))
        g.apply_edges(fn.e_mul_v(out_name, accum_name, out_name))
        g.ndata.pop(accum_name)
        grad_score = g.edata.pop(grad_score_name) - g.edata.pop(out_name)
        return grad_score


def edge_softmax(graph, logits):

After Change

        sds_sum = sds.dst_sum()  // type dgl.NData
        grad_score = sds - sds * sds_sum  // multiple expressions
        
        g = self.g.local_var()
        out, = self.saved_tensors  // pylint: disable=access-member-before-definition, unpacking-non-sequence
        // clear saved tensors explicitly
        self.saved_tensors = None
        g.edata["out"] = out
        g.edata["grad_score"] = out * grad_out
        g.update_all(fn.copy_e("grad_score", "m"), fn.sum("m", "accum"))
        g.apply_edges(fn.e_mul_v("out", "accum", "out"))
        grad_score = g.edata["grad_score"] - g.edata["out"]
        return grad_score

def edge_softmax(graph, logits):

Instances