Add a batch of weighted updates into a histogram and return the new histogram. JAX arrays cannot be changed in place, so the result has to be built as a new array.
counts holds the starting value of each of the K bins. Update number n adds weights[n] to the bin whose number is indices[n], so for every bin k
resultk=countsk+n:indicesn=k∑weightsnwhere the sum runs over every update n whose index equals k. The same bin number can appear many times in indices, and every one of those updates must be counted. An empty indices array leaves the counts as they are.
Only the indices 0 to K - 1 name a bin. Some updates carry an index outside that range, for example -1 for an item that an earlier step discarded, or K for an item past the last bin. Every such update, negative or too large, must be ignored and changes no bin.
Your function is called under jax.jit, so its inputs are traced arrays whose values are unknown while the function is being compiled.
Return a new float32 JAX array of shape (K,).
Input: counts = [0.0,0.0,0.0,0.0], indices = [1,3,1,0], weights = [1.0,1.0,1.0,1.0]
Output: [1.0,2.0,0.0,1.0]
Explanation: Bin 1 appears twice in indices, so it receives two weights of 1. Bins 0 and 3 receive one weight each, and bin 2 receives none.
Input: counts = [0.5,0.0,2.0], indices = [2,-1,2,3,0], weights = [0.25,0.5,1.0,4.0,-1.0]
Output: [-0.5,0.0,3.25]
Explanation: The indices -1 and 3 lie outside the three bins, so their weights 0.5 and 4.0 are ignored. Bin 2 receives 0.25 and 1.0 on top of its starting 2.0, and bin 0 receives -1.0.
Input: counts = [1.0,2.0,3.0], indices = [], weights = []
Output: [1.0,2.0,3.0]
counts.at[indices].add(weights) adds every weight into its bin, and repeated indices accumulate.
The .at methods take mode="drop", which skips indices past the end but still wraps negative ones.
jnp.where(valid, weights, 0.0) zeroes the weights you want to ignore.
Accepts: array
Accepts: array
Accepts: array
Add a batch of weighted updates into a histogram and return the new histogram. JAX arrays cannot be changed in place, so the result has to be built as a new array.
counts holds the starting value of each of the K bins. Update number n adds weights[n] to the bin whose number is indices[n], so for every bin k
resultk=countsk+n:indicesn=k∑weightsnwhere the sum runs over every update n whose index equals k. The same bin number can appear many times in indices, and every one of those updates must be counted. An empty indices array leaves the counts as they are.
Only the indices 0 to K - 1 name a bin. Some updates carry an index outside that range, for example -1 for an item that an earlier step discarded, or K for an item past the last bin. Every such update, negative or too large, must be ignored and changes no bin.
Your function is called under jax.jit, so its inputs are traced arrays whose values are unknown while the function is being compiled.
Return a new float32 JAX array of shape (K,).
Input: counts = [0.0,0.0,0.0,0.0], indices = [1,3,1,0], weights = [1.0,1.0,1.0,1.0]
Output: [1.0,2.0,0.0,1.0]
Explanation: Bin 1 appears twice in indices, so it receives two weights of 1. Bins 0 and 3 receive one weight each, and bin 2 receives none.
Input: counts = [0.5,0.0,2.0], indices = [2,-1,2,3,0], weights = [0.25,0.5,1.0,4.0,-1.0]
Output: [-0.5,0.0,3.25]
Explanation: The indices -1 and 3 lie outside the three bins, so their weights 0.5 and 4.0 are ignored. Bin 2 receives 0.25 and 1.0 on top of its starting 2.0, and bin 0 receives -1.0.
Input: counts = [1.0,2.0,3.0], indices = [], weights = []
Output: [1.0,2.0,3.0]
counts.at[indices].add(weights) adds every weight into its bin, and repeated indices accumulate.
The .at methods take mode="drop", which skips indices past the end but still wraps negative ones.
jnp.where(valid, weights, 0.0) zeroes the weights you want to ignore.
Accepts: array
Accepts: array
Accepts: array
