PyTorch BinCount

PyTorch 0.4.1 now supports bincount both for CPU and GPU. Therefore, this package is no longer needed and will not be updated.

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.

Installation

Ensure that at least PyTorch 0.4.1 is installed and verify that cuda/bin and cuda/include are in your $PATH and $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/include:...

Then run:

pip install torch-bincount

If you are running into any installation problems, please create an issue. Be sure to import torch first before using this package to resolve symbols the dynamic linker must see.

Usage

torch_bincount.bincount(src, size=None) -> LongTensor

Counts the number of occurrences of each value in a non-negative tensor.

Parameters

src (Tensor) - The input tensor.
size (int, optional) - The maximum number of bins for the output array. (default: None)

Returns

out (LongTensor) - The result of binning the input tensor.

Example

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])

Running tests

python setup.py test

rusty1s / pytorch_bincount