tensorflow
tf.unsortedSegmentSum does not aggregate all indices?
- 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 only 88804 elements are counted?
I am overcoming this issue by using following code. The actual distinct segment number is just size*size, but I have to supply a larger buffer to solve it.
const ONES_TENSOR = tf.ones([indexTensor.size], 'float32');
// using tf.unsortedSegmentSum to calculate the sum of the ones based on the indices in indexTensor
const countTensor = tf.unsortedSegmentSum(ONES_TENSOR, indexTensor, indexTensor.size);
indexTensor.dispose();
// Possible bug in TensorflowJS : https://github.com/tensorflow/tfjs/issues/8174
const slicedTensor = tf.slice(countTensor, 0, size * size);
countTensor.dispose();
Hi, @wangjia184
I apologize for the delayed response, Good to hear that your issue has been resolved by supplying larger buffer size. I was trying to replicate the same behavior from my end with webgpu backend in Tensorflow.js api page itself and I'm getting below output at the moment I'm not sure what values you are passing so I tried with size=1000 & indexTensor is equal to const indexTensor = tf.ones([size * size], 'int32');
If I have missed something here please let me know. Thank you.
Thanks @gaikwadrahul8 , please find Pen.io here: https://codepen.io/wangjia184/pen/NWJoXdq and watch console output
<html>
<head>
<!-- Import @tensorflow/tfjs or @tensorflow/tfjs-core -->
<script src="https://cdn.jsdelivr.net/npm/@tensorflow/tfjs/dist/tf.min.js"> </script>
<!-- Add the WebGPU backend to the global backend registry -->
<script src="https://cdn.jsdelivr.net/npm/@tensorflow/tfjs-backend-webgpu/dist/tf-backend-webgpu.js"></script>
</head>
<body onload="init()">
<script>
function init() {
// Set the backend to WebGPU and wait for the module to be ready
tf.setBackend('webgpu').then( async () => {
console.log('WebGPU is ready');
const TOTAL_SAMPLES = 1000000;
const shape = [TOTAL_SAMPLES];
// generate noises
const guassianTensor = tf.randomNormal(shape, 0/*mean*/, 1/*stddev*/, 'float32');
console.log("guassianTensor shape =", guassianTensor.shape, "; dtype =", guassianTensor.dtype);
guassianTensor.print();
const segmentNums = 256*256;
const increment = 2 / (segmentNums - 1.0); // range from [-1, +1]
const edges = Array.from({ length: segmentNums - 2 }, (_, idx) => -1 + (idx + 1) * increment);
const seqTensor = tf.tensor(edges);
console.log("seqTensor shape =", seqTensor.shape, "; dtype =", seqTensor.dtype);
seqTensor.print();
const discretizedTensor = tf.lowerBound(seqTensor, guassianTensor);
seqTensor.dispose();
guassianTensor.dispose();
const indexTensor = discretizedTensor.reshape([TOTAL_SAMPLES]);
discretizedTensor.dispose();
console.log("indexTensor shape =", indexTensor.shape, "; dtype =", indexTensor.dtype);
indexTensor.print();
const indexBuffer = await indexTensor.array();
console.log("indexBuffer valid element number ="
, indexBuffer.map((v) => (v >= 0 && v < segmentNums) ? 1 : 0).reduce((a, b) => a + b)
);
const onesTensor = tf.ones([segmentNums], 'float32');
console.log("onesTensor shape =", onesTensor.shape, "; dtype =", onesTensor.dtype);
const countTensor = tf.unsortedSegmentSum(onesTensor, indexTensor, segmentNums);
console.log("countTensor shape =", countTensor.shape, "; dtype =", countTensor.dtype);
const data = await countTensor.array();
console.log("total =", data.reduce((a, b) => a + b));
console.log(data);
indexTensor.dispose();
onesTensor.dispose();
countTensor.dispose();
});
}
</script>
</body>
</html>
Hi, @wangjia184
I apologize for the delayed response and I tried to replicate the same behaviour from my end and I'm also getting the same result which you mentioned in your above comment so we'll have to dig more into this, thank you for bringing this issue to our attention
For reference I have added output log below :
Thank you for your understanding and patience.