Mass matrix multiplication in tensorflow

loadScript.js

function _loadScript(path) {
  var script = document.createElement('script');
  script.type = 'text/javascript';
  script.src = path;
  document.head.appendChild(script);
}

tensorflow.js

_loadScript('https://unpkg.com/@tensorflow/tfjs');
// WebGL backend crashes
tf.setBackend('cpu');
// Create a transformation matrix
const A = tf.tensor2d([1, 0, 0, 0, 1, 0, 0, 0, 1], [3, 3]);
const pixelCreation = performance.now()
// Create a list of point representing the pixel of the camera
const incrVectorX = [0.5, 0, 0]; // shape (3,)
const incrVectorY = [0, 0.5, 0]; // shape (3,)
const xindices = tf.range(0, 512).expandDims(1); // shape (512, 1)
const yindices = tf.range(0, 512).expandDims(1); // shape (512, 1)
const b = (xindices.mul(incrVectorX).expandDims(0).add(yindices.mul(incrVectorY).expandDims(1))).reshape([512 * 512, 3])
console.log(`stack created in ${(performance.now() - pixelCreation).toFixed(2)} ms`);

const multiplication = performance.now();
// Multiply the matrix by all those vectors
const result = A.matMul(b.transpose());
console.log(`Multiplication done in ${(performance.now() - multiplication).toFixed(2)} ms`);