a27d95173d560a536480ef86ce18640bb6ee95b8,torch_geometric/utils/softmax.py,,softmax,#Any#Any#Any#Any#,9

Before Change

    :rtype: :class:`Tensor`
    
    out = src
    if src.numel() > 0:
        out = out - src.max()
    out = out.exp()

    if ptr is not None:
        out_sum = gather_csr(segment_csr(out, ptr, reduce="sum"), ptr)

After Change

    :rtype: :class:`Tensor`
    
    if ptr is not None:
        src_max = gather_csr(segment_csr(src, ptr, reduce="max"), ptr)
        out = (src - src_max).exp()
        out_sum = gather_csr(segment_csr(out, ptr, reduce="sum"), ptr)
    elif index is not None:
        N = maybe_num_nodes(index, num_nodes)
        src_max = scatter(src, index, dim=0, dim_size=N, reduce="max")[index]

Instances