min.py 3.79 KB
Newer Older
rusty1s's avatar
rusty1s committed
1
import torch
rusty1s's avatar
rusty1s committed
2
3
from torch.autograd import Function

rusty1s's avatar
rusty1s committed
4
5
from torch_scatter.utils.ext import get_func
from torch_scatter.utils.gen import gen
rusty1s's avatar
rusty1s committed
6
7
8
9
10
11
12


class ScatterMin(Function):
    @staticmethod
    def forward(ctx, out, src, index, dim):
        arg = index.new_full(out.size(), -1)
        func = get_func('scatter_min', src)
rusty1s's avatar
rusty1s committed
13
        func(src, index, out, arg, dim)
rusty1s's avatar
rusty1s committed
14
15
16
17
18
19
20
21
22

        ctx.mark_dirty(out)
        ctx.dim = dim
        ctx.save_for_backward(index, arg)

        return out, arg

    @staticmethod
    def backward(ctx, grad_out, grad_arg):
23
        index, arg = ctx.saved_tensors
rusty1s's avatar
rusty1s committed
24
25
26

        grad_src = None
        if ctx.needs_input_grad[1]:
rusty1s's avatar
rusty1s committed
27
28
29
30
31
            size = list(index.size())
            size[ctx.dim] += 1
            grad_src = grad_out.new_zeros(size)
            grad_src.scatter_(ctx.dim, arg.detach() + 1, grad_out)
            grad_src = grad_src.narrow(ctx.dim, 1, index.size(ctx.dim))
rusty1s's avatar
rusty1s committed
32
33
34
35

        return None, grad_src, None, None


rusty1s's avatar
rusty1s committed
36
def scatter_min(src, index, dim=-1, out=None, dim_size=None, fill_value=None):
rusty1s's avatar
rusty1s committed
37
38
39
40
41
42
43
44
45
46
47
    r"""
    |

    .. image:: https://raw.githubusercontent.com/rusty1s/pytorch_scatter/
            master/docs/source/_figures/min.svg?sanitize=true
        :align: center
        :width: 400px

    |

    Minimizes all values from the :attr:`src` tensor into :attr:`out` at the
rusty1s's avatar
typo  
rusty1s committed
48
    indices specified in the :attr:`index` tensor along a given axis
rusty1s's avatar
rusty1s committed
49
    :attr:`dim`.If multiple indices reference the same location, their
50
    **contributions minimize** (`cf.` :meth:`~torch_scatter.scatter_add`).
rusty1s's avatar
rusty1s committed
51
    The second return tensor contains index location in :attr:`src` of each
52
    minimum value (known as argmin).
rusty1s's avatar
rusty1s committed
53
54
55
56
57
58

    For one-dimensional tensors, the operation computes

    .. math::
        \mathrm{out}_i = \min(\mathrm{out}_i, \min_j(\mathrm{src}_j))

rusty1s's avatar
rusty1s committed
59
    where :math:`\min_j` is over :math:`j` such that
rusty1s's avatar
rusty1s committed
60
61
62
63
64
65
66
67
68
69
70
71
72
    :math:`\mathrm{index}_j = i`.

    Args:
        src (Tensor): The source tensor.
        index (LongTensor): The indices of elements to scatter.
        dim (int, optional): The axis along which to index.
            (default: :obj:`-1`)
        out (Tensor, optional): The destination tensor. (default: :obj:`None`)
        dim_size (int, optional): If :attr:`out` is not given, automatically
            create output with size :attr:`dim_size` at dimension :attr:`dim`.
            If :attr:`dim_size` is not given, a minimal sized output tensor is
            returned. (default: :obj:`None`)
        fill_value (int, optional): If :attr:`out` is not given, automatically
rusty1s's avatar
rusty1s committed
73
74
75
76
77
            fill output tensor with :attr:`fill_value`. (default: :obj:`None`)
        fill_value (int, optional): If :attr:`out` is not given, automatically
            fill output tensor with :attr:`fill_value`. If set to :obj:`None`,
            the output tensor is filled with the greatest possible value of
            :obj:`src.dtype`. (default: :obj:`None`)
rusty1s's avatar
rusty1s committed
78
79
80
81
82
83
84
85
86

    :rtype: (:class:`Tensor`, :class:`LongTensor`)

    .. testsetup::

        import torch

    .. testcode::

rusty1s's avatar
rusty1s committed
87
88
        from torch_scatter import scatter_min

rusty1s's avatar
rusty1s committed
89
        src = torch.Tensor([[-2, 0, -1, -4, -3], [0, -2, -1, -3, -4]])
rusty1s's avatar
rusty1s committed
90
91
        index = torch.tensor([[ 4, 5,  4,  2,  3], [0,  0,  2,  2,  1]])
        out = src.new_zeros((2, 6))
rusty1s's avatar
rusty1s committed
92

93
        out, argmin = scatter_min(src, index, out=out)
rusty1s's avatar
rusty1s committed
94

rusty1s's avatar
rusty1s committed
95
        print(out)
96
        print(argmin)
rusty1s's avatar
rusty1s committed
97
98
99

    .. testoutput::

rusty1s's avatar
rusty1s committed
100
101
       tensor([[ 0.,  0., -4., -3., -2.,  0.],
               [-2., -4., -3.,  0.,  0.,  0.]])
rusty1s's avatar
rusty1s committed
102
103
       tensor([[-1, -1,  3,  4,  0,  1],
               [ 1,  4,  3, -1, -1, -1]])
rusty1s's avatar
rusty1s committed
104
    """
rusty1s's avatar
rusty1s committed
105
106
107
    if fill_value is None:
        op = torch.finfo if torch.is_floating_point(src) else torch.iinfo
        fill_value = op(src.dtype).max
rusty1s's avatar
rusty1s committed
108
    src, out, index, dim = gen(src, index, dim, out, dim_size, fill_value)
109
    if src.size(dim) == 0:  # pragma: no cover
110
        return out, index.new_full(out.size(), -1)
rusty1s's avatar
rusty1s committed
111
    return ScatterMin.apply(out, src, index, dim)