How to activate webgpu backend on Tensorflow.JS

Question

Current release on Github of Tensorflow.js includes a WEBGPU backend, but when running tf.getBackend() I get webgl even with Chrome Canary with WEBGPU enabled. Also running tf.backend() doesn't seem to indicate that WEBGPU is there.

Answer 1

Which tf package are you importing? The WebGPU backend is not bundled with tfjs by default - you need to import it specifically, like this:

import * as tf from '@tensorflow/tfjs-backend-webgpu';

const init = async () => {
    await tf.ready();

    // Now we can create tensors and run ops.
    tf.matMul(a, b).print(); 
};

init();

Note that the backend is asynchronous - you have to call await tf.ready() before doing anything with it.

Hope that helps.