"vscode:/vscode.git/clone" did not exist on "8ae8c9fa8ce19120a2ec44d0b799727bea1edad7"
min.py 3.29 KB
Newer Older
rusty1s's avatar
rusty1s committed
1
2
3
4
5
from .scatter import scatter
from .utils import gen_filled_tensor, gen_output


def scatter_min_(output, index, input, dim=0):
rusty1s's avatar
rusty1s committed
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
    r"""Minimizes all values from the :attr:`input` tensor into :attr:`output`
    at the indices specified in the :attr:`index` tensor along an given axis
    :attr:`dim`. If multiple indices reference the same location, their
    **contributions minimize** (`cf.` :meth:`~torch_scatter.scatter_add_`).
    The second return value is the index location in :attr:`input` of each
    minimum value found (argmin).

    For one-dimensional tensors, the operation computes

    .. math::
        \mathrm{output}_i = \min(\mathrm{output}_i, \min_j(\mathrm{input}_j))

    where min is over :math:`j` such that :math:`\mathrm{index}_j = i`.

    Args:
        output (Tensor): The destination tensor
        index (LongTensor): The indices of elements to scatter
        input (Tensor): The source tensor
        dim (int, optional): The axis along which to index

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

    .. testsetup::

        import torch

    .. testcode::

        from torch_scatter import scatter_min_
        input =     torch.Tensor([[-2, 0, -1, -4, -3], [0, -2, -1, -3, -4]])
        index = torch.LongTensor([[ 4, 5,  4,  2,  3], [0,  0,  2,  2,  1]])
        output = torch.zeros(2, 6)
        output = scatter_min_(output, index, input, dim=1)
        print(output)

    .. testoutput::

       (
        0  0 -4 -3 -2  0
       -2 -4 -3  0  0  0
       [torch.FloatTensor of size 2x6]
       ,
       -1 -1  3  4  0  1
        1  4  3 -1 -1 -1
       [torch.LongTensor of size 2x6]
       )
    """
rusty1s's avatar
rusty1s committed
53
54
55
56
    arg_output = gen_filled_tensor(index, output.size(), fill_value=-1)
    return scatter('min', dim, output, index, input, arg_output)


rusty1s's avatar
rename  
rusty1s committed
57
def scatter_min(index, input, dim=0, size=None, fill_value=0):
rusty1s's avatar
rusty1s committed
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
    r"""Minimizes all values from the :attr:`input` tensor at the indices
    specified in the :attr:`index` tensor along an given axis :attr:`dim`
    (`cf.` :meth:`~torch_scatter.scatter_min_` and
    :meth:`~torch_scatter.scatter_add`).

    For one-dimensional tensors, the operation computes

    .. math::
        \mathrm{output}_i = \min(\mathrm{fill\_value},
        \min_j(\mathrm{input}_j))

    where min is over :math:`j` such that :math:`\mathrm{index}_j = i`.

    Args:
        index (LongTensor): The indices of elements to scatter
        input (Tensor): The source tensor
        dim (int, optional): The axis along which to index
        size (int, optional): Output size at dimension :attr:`dim`
        fill_value (int, optional): Initial filling of output tensor

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

    .. testsetup::

        import torch

    .. testcode::

        from torch_scatter import scatter_min
        input =     torch.Tensor([[-2, 0, -1, -4, -3], [0, -2, -1, -3, -4]])
        index = torch.LongTensor([[ 4, 5,  4,  2,  3], [0,  0,  2,  2,  1]])
        output = scatter_min(index, input, dim=1)
        print(output)

    .. testoutput::

       (
        0  0 -4 -3 -2  0
       -2 -4 -3  0  0  0
       [torch.FloatTensor of size 2x6]
       ,
       -1 -1  3  4  0  1
        1  4  3 -1 -1 -1
       [torch.LongTensor of size 2x6]
       )
    """
rusty1s's avatar
rename  
rusty1s committed
104
    output = gen_output(index, input, dim, size, fill_value)
rusty1s's avatar
rusty1s committed
105
    return scatter_min_(output, index, input, dim)