Tensorflow.js tf.customGrad() Function

// 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();

// 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();