Tensorflow.js tf.customGrad() Function

Tensorflow.js is an open-source library developed by Google for running machine learning models and deep learning neural networks in the browser or node environment.

The tf.customGrad() function is used to return the gradient of a specified custom function “f”. Here the custom function gives {value: Tensor, gradFunc: (dy, saved) → Tensor[]}, where gradFunc gives the custom gradients of the input function f in respect of its inputs.

Syntax:

tf.customGrad(f)

Parameters: This function accepts a parameter which is illustrated below:

• f: It is the specified custom function.

Return Value: This function returns the gradient of a specified custom function “f”

Example 1:

Javascript

// Importing the tensorflow.js library
import * as tf from "@tensorflow/tfjs"

// Initializing a custom function f
const f = (a, save) => {

   // Saving a for its availability later for the gradient
   save([a]);

   // Overriding gradient of a^2
   return {
     value: a.square(),

     // Here "saved.a" pointing to "a" which
     // have been saved above
     gradFunc: (dy, saved) => [dy.mul(saved[0].abs())]
   };
}

// Calling the .customGrad() function
// over the above specified custom function f
const customOp = tf.customGrad(f);

// Initializing a 1D Tensor of some values
const a = tf.tensor1d([0, -1, 1, 2]);

// Getting the gradient of above function
// f for the above specified Tensor values
const da = tf.grad(a => customOp(a));

// Printing the custom function "f" for the
// above specified Tensor "a"
console.log(`f(a):`);
customOp(a).print();

// Printing the gradient of the function "f" for the
// above specified Tensor "a"
console.log(`f'(a):`);
da(a).print();

Output:

f(a):
Tensor
   [0, 1, 1, 4]
f'(a):
Tensor
   [0, 1, 1, 2]

Example 2:

Javascript

// Importing the tensorflow.js library
import * as tf from "@tensorflow/tfjs"

// Calling the .customGrad() function
// with the custom function "f" as
// it's parameter
const customOp = tf.customGrad(

// Initializing a custom function f
(a, save) => {

   // Saving a for its availability later for the gradient
   save([a]);

   // Overriding gradient of a^3
   return {
     value: a.pow(tf.scalar(3, 'int32')),

     // Here "saved.a" pointing to "a" which
     // have been saved above
     gradFunc: (dy, saved) => [dy.mul(saved[0].abs())]
   };
}
);

// Initializing a 1D Tensor of some values
const a = tf.tensor1d([0, -1, 2, -2, 0.3]);

// Getting the gradient of above function
// f for the above specified Tensor values
const da = tf.grad(a => customOp(a));

// Printing the custom function "f" for the
// above specified Tensor "a"
console.log(`f(a):`);
customOp(a).print();

// Printing the gradient of the function "f" for the
// above specified Tensor "a"
console.log(`f'(a):`);
da(a).print();

Output:

f(a):
Tensor
   [0, -1, 8, -8, 0.027]
f'(a):
Tensor
   [0, 1, 2, 2, 0.3]

Reference:https://js.tensorflow.org/api/latest/#customGrad