TensorFlow.js

Sections


Creating custom WebGL operations

To define custom WebGL operations, all we have to do is create an object that implements tf.webgl.GPGPUProgram.

This interface is defined as:

interface GPGPUProgram {
  variableNames: string[];
  outputShape: number[];
  userCode: string;
  supportsBroadcasting?: boolean;
}

For a contrived example, lets implement an operation that computes f(x) = x * x + x.

The GLSL code for this would be:

void main() {
    float x = getXAtOutCoords();
    float value = x * x + x;
    setOutput(value);
}

where getXAtOutCoords and setOutput are provided by Tensorflow.js to the shader. Note that the main function is called for each value in the output tensor.

The full GPGPUProgram definition would be:

const squareAndAddKernel = inputShape => ({
  variableNames: ['X'],
  outputShape: inputShape.slice(),
  userCode: `
    void main() {
        float x = getXAtOutCoords();
        float value = x * x + x;
        setOutput(value);
      }
  `
})

To run this op, you would use tf.ENV.backend.compileAndRun(program: GPGPUProgram, inputs: tf.Tensor[]): tf.Tensor. Note that this will be undefined if the backend isn't the webgl backend.

const x = tf.tensor([1, 2, 3, 4]);
const program = squareAndAddKernel(x.shape);

const result = tf.ENV.backend.compileAndRun(program, [x]);

However, we probably also want to define the gradients for this op, so that gradients can be backpropagated through it.

To do this, we use tf.customGrad.

const squareAndAddBackpropKernel = inputShape => ({
  variableNames: ['X'],
  outputShape: inputShape.slice(),
  userCode: `
    void main() {
      float x = getXAtOutCoords();
      float value = 2.0 * x + 1.0;
      setOutput(value);
    }
  `
});


const squareAndAdd = tf.customGrad(x => {
  const backend = tf.ENV.backend;
  const program = squareAndAddKernel(x.shape);
  const backpropProgram = squareAndAddBackpropKernel(x.shape);

  const value = backend.compileAndRun(program, [x]);

  const gradFunc = dy =>
      [backend.compileAndRun(backpropProgram, [x]).mul(dy)];
  return {value, gradFunc};
});

We can then use this as:

const x = tf.tensor([1, 2, 3, 4]);

const value = squareAndAdd(x);

const grads = tf.grad(x => squareAndAdd(x));
const dx = grads(x);

// value == [2, 6, 12, 20]
// dx == [3, 5, 7, 9]

Or more concisely:

const {value, grad} = tf.valueAndGrad(squareAndAdd)(x);

GLSL functions generated by Tensorflow.js

Tensorflow.js generates functions you can use to read from the input tensors and write to the output tensor, as well as additional numeric utility functions. These are prepended to your code by the Shader Compiler.