Optimal way for calc histogram in tensorflow.js - tensorflow.js

Question

I need to calc histogram from big 1D tensor with tensorflow.js. I found way to do it, but its not optimal because of hi memory cosumtion in oneHot operation. My code example below:
for (let index = 0; index < 50; index++) { //repeat some times to ensure no memory leaks
const hist = getnormalHist(index);
const histArray = (await hist.array()) as number[];
const values = histArray.map((v, i) => ({ index: i, value: v }));
await tfvis.render.barchart({ name: "hist" }, values);
}
function getnormalHist(seed: number) {
return tf.tidy(() => {
const rand = tf
.randomNormal([100000], 0, 100, "int32", seed) //Generates long array of normal distributed randoms
.add(500)
.toInt();
const oneHot = rand.oneHot(1000, 1, 0); // convert to oneHot makes it x1000 bigger
const hist = oneHot.transpose().sum(1); // finally get histogram tensor
return hist;
});
}
I need to make this code faster and with less memory consumption, but I don't understand how.

Answer 1

tf.randomNormal([100000], 0, 100, "int32", seed) //Generates long array of normal distributed randoms
.add(500)
.toInt()
The above operation generates a sample distribution with mean around 500. It can be generated using mean 500 instead of 0. The add operation is not necessary for that. Additionnally it is redundant to use toInt since the dtype is already int32. Thus it can be simplified this way:
tf.randomNormal([100000], 500, 100, "int32", seed)
It is redundant to specify 0 and 1 onvalue and offvalue as they are default value.
rand.oneHot(1000)
There is no need to transpose before computing the sum for each indice. Computing the sum over the 0 axis will count each index. An intermediate matrix of size 100000 * 1000 will no longer be used.
oneHot.sum(0)
As a summary, getNormalHist will look as the following:
function getnormalHist(seed: number) {
return tf.tidy(() => {
const rand = tf
.randomNormal([100000], 500, 100, "int32", seed) //Generates long array of normal distributed randoms
const oneHot = rand.oneHot(1000); // convert to oneHot makes it x1000 bigger
const hist = oneHot.sum(0); // finally get histogram tensor
return hist;
});
}

Answer 2

My solution is create custom webGL operation.
class HistProgram {
variableNames = ["X"];
outputShape: number[];
userCode: string;
constructor(numIndices: number, counts: number) {
this.outputShape = [numIndices];
this.userCode = `
void main() {
int ch = getOutputCoords();
int c = 0;
for (int i = 0; i < ${counts}; i++){
if(ch == int(getX(i))){
c++;
}
}
setOutput(float(c));
}
`;
}
}
Now i can use it in this way:
const counts = 100000;
const channels = 1000;
const rand = tf.randomNormal([counts], 500, 100, "int32")
const prog = new HistProgram(channels, counts);
const hist = await tf.backend().compileAndRun(prog, [rand]).array();

Answer 3

Now the best solution that i found is:
const counts = 100000;
const channels = 1000;
const rand = tf.randomNormal([counts], 500, 100, "int32");
const hist = tf.sparseToDense(rand, tf.onesLike(rand), [channels], 0);