README.md 685 Bytes
Newer Older
rusty1s's avatar
rusty1s committed
1
# PyTorch Scatter
rusty1s's avatar
rusty1s committed
2

rusty1s's avatar
rusty1s committed
3
<p align="center">
rusty1s's avatar
smaller  
rusty1s committed
4
  <img width="50%" src="https://raw.githubusercontent.com/rusty1s/pytorch_scatter/master/docs/source/_figures/add.svg?sanitize=true" />
rusty1s's avatar
rusty1s committed
5
6
7
8
</p>

--------------------------------------------------------------------------------

rusty1s's avatar
rusty1s committed
9
10
11
12
13
14
15
```sh
pip install torch_scatter
```

```py
from torch_scatter import scatter_max

rusty1s's avatar
typos  
rusty1s committed
16
17
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]])
rusty1s's avatar
rusty1s committed
18

rusty1s's avatar
typos  
rusty1s committed
19
max, argmax = scatter_max_(index, input, dim=1)
rusty1s's avatar
typo  
rusty1s committed
20
```
rusty1s's avatar
rusty1s committed
21

rusty1s's avatar
typo  
rusty1s committed
22
```
rusty1s's avatar
typos  
rusty1s committed
23
24
25
26
print(max)
 0  0  4  3  2  0
 2  4  3  0  0  0
[torch.FloatTensor of size 2x6]
rusty1s's avatar
rusty1s committed
27

rusty1s's avatar
typos  
rusty1s committed
28
29
30
31
32
print(argmax)
-1 -1  3  4  0  1
 1  4  3 -1 -1 -1
[torch.LongTensor of size 2x6]
```