This package consists of a small extension library of a highly optimized
bincount operation for the use in PyTorch, which is missing in the main package.
The operation works on varying data types and is implemented both for CPU and GPU.
Ensure that at least PyTorch 0.4.1 is installed and verify that
cuda/install are in your
$CPATH respectively, e.g.:
$ python -c "import torch; print(torch.__version__)" >>> 0.4.1 $ echo $PATH >>> /usr/local/cuda/bin:... $ echo $CPATH >>> /usr/local/cuda/install:...
pip install torch-scatter torch-bincount
If you are running into any installation problems, please create an issue.
torch_bincount.bincount(src, size=None) -> LongTensor
Counts the number of occurrences of each value in a non-negative tensor.
- src (Tensor) - The input tensor.
size (int, optional) - The maximum number of bins for the output array. (default:
- out (LongTensor) - The result of binning the input tensor.
import torch from torch_bincount import bincount src = torch.tensor([2, 1, 1, 2, 4, 4, 2]) out = bincount(src)
print(out) tensor([ 0, 2, 3, 0, 2])
python setup.py test