- Tensorflow JS : version 4.16.0
- Backend : webgpu
// https://js.tensorflow.org/api/4.17.0/#unsortedSegmentSum
const onesTensor = tf.ones([size * size], 'float32');
const countTensor = tf.unsortedSegmentSum(onesTensor, indexTensor, size * size);
console.log("onesTensor shape =", onesTensor.shape, "; dtype =", onesTensor.dtype);
console.log("indexTensor shape =", indexTensor.shape, "; dtype =", indexTensor.dtype);
console.log("countTensor shape =", countTensor.shape, "; dtype =", countTensor.dtype);
console.log("size * size =", size * size);
const indexBuffer = indexTensor.arraySync();
console.log("indexBuffer valid element number ="
, indexBuffer.map((v) => (v >= 0 && v < size * size) ? 1 : 0).reduce((a, b) => a + b)
);
const data = countTensor.arraySync();
console.log("total =", data.reduce((a, b) => a + b));
I have the above code using tf.unsortedSegmentSum to calculate the sum of the ones based on the indices in indexTensor.
Here is output from console.
onesTensor shape = [88804] ; dtype = float32
indexTensor shape = [1000000] ; dtype = int32
countTensor shape = [88804] ; dtype = float32
size * size = 88804
indexBuffer valid element number = 1000000
total = 88804
As you can see, there are 1000000 indices in indexBuffer but there are only 88804 indices aggregated. Actually those 88804 indices are aggregated correctly, I can see duplicates are counted, I can see the graph if I draw in heatmap. But there are 1000000 valid indices and total should be 1000000.
Why are only 88804 indices counted?