So I tried pasting the 512x512 matrix multiplication code in a snippet. It worked but I noticed that GPU was over 4x slower than CPU (using the performance tool). I figured that using mode: webgl was way faster than webgl2 (and cpu). That was on Ubuntu 18.04 with a GTX 1070.

I added an Observable notebook: https://observablehq.com/@benoit9/gpu-js-matrix-multiply-webgl2-issue