tf.unsortedSegmentSum does not aggregate all indices?

25 views Asked by At
  • 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?

0

There are 0 answers