jax.numpy.bincount¶

jax.numpy.bincount(x, weights=None, minlength=0, *, length=None)[source]¶

Count number of occurrences of each value in array of non-negative ints.

LAX-backend implementation of bincount(). Jax adds the optional length parameter which specifies the output length, and defaults to x.max() + 1. It must be specified for bincount to be compilable. Values larger than the specified length will be discarded.

Additionally, while np.bincount raises an error if the input array contains negative values, jax.numpy.bincount treats negative values as zero.

Original docstring below.

bincount(x, weights=None, minlength=0)

The number of bins (of size 1) is one larger than the largest value in x. If minlength is specified, there will be at least this number of bins in the output array (though it will be longer if necessary, depending on the contents of x). Each bin gives the number of occurrences of its index value in x. If weights is specified the input array is weighted by it, i.e. if a value n is found at position i, out[n] += weight[i] instead of out[n] += 1.
Returns
outndarray of ints
The result of binning the input array. The length of out is equal to np.amax(x)+1.

ValueError
If the input is not 1-dimensional, or contains elements with negative values, or if minlength is negative.

TypeError
If the type of the input is float or complex.

histogram, digitize, unique
>>> np.bincount(np.arange(5))
array([1, 1, 1, 1, 1])
>>> np.bincount(np.array([0, 1, 1, 3, 2, 1, 7]))
array([1, 3, 1, 1, 0, 0, 0, 1])
>>> x = np.array([0, 1, 1, 3, 2, 1, 7, 23])
>>> np.bincount(x).size == np.amax(x)+1
True
The input array needs to be of integer dtype, otherwise a TypeError is raised:
>>> np.bincount(np.arange(5, dtype=float))
Traceback (most recent call last):
  ...
TypeError: Cannot cast array data from dtype('float64') to dtype('int64')
according to the rule 'safe'
A possible use of bincount is to perform sums over variable-size chunks of an array, using the weights keyword.
>>> w = np.array([0.3, 0.5, 0.2, 0.7, 1., -0.6]) # weights
>>> x = np.array([0, 1, 1, 2, 2, 2])
>>> np.bincount(x,  weights=w)
array([ 0.3,  0.7,  1.1])