Method customGrad

tfjs-core/src/engine.ts:1157–1216

(f: CustomGradientFunc<T>)

  }

  customGrad<T extends Tensor>(f: CustomGradientFunc<T>):
      (...args: Array<Tensor|GradSaveFunc>) => T {
    util.assert(
        util.isFunction(f),
        () => 'The f passed in customGrad(f) must be a function.');
    return (...inputs: Tensor[]): T => {
      util.assert(
          inputs.every(t => t instanceof Tensor),
          () => 'The args passed in customGrad(f)(x1, x2,...) must all be ' +
              'tensors');

      let res: {
        value: T,
        gradFunc: (dy: T, saved: Tensor[]) => Tensor | Tensor[],
      };
      const inputMap: NamedTensorMap = {};
      inputs.forEach((input, i) => {
        inputMap[i] = input;
      });

      const forwardFunc: ForwardFunc<T> = (_, save) => {
        res = f(...[...inputs, save]);
        util.assert(
            res.value instanceof Tensor,
            () => 'The function f passed in customGrad(f) must return an ' +
                'object where `obj.value` is a tensor');
        util.assert(
            util.isFunction(res.gradFunc),
            () => 'The function f passed in customGrad(f) must return an ' +
                'object where `obj.gradFunc` is a function.');
        return res.value;
      };

      const backwardsFunc = (dy: T, saved: Tensor[]) => {
        const gradRes = res.gradFunc(dy, saved);
        const grads: Tensor[] = Array.isArray(gradRes) ? gradRes : [gradRes];
        util.assert(
            grads.length === inputs.length,
            () => 'The function f passed in customGrad(f) must return an ' +
                'object where `obj.gradFunc` is a function that returns ' +
                'the same number of tensors as inputs passed to f(...).');
        util.assert(
            grads.every(t => t instanceof Tensor),
            () => 'The function f passed in customGrad(f) must return an ' +
                'object where `obj.gradFunc` is a function that returns ' +
                'a list of only tensors.');
        const gradMap: {[key: string]: () => Tensor} = {};
        grads.forEach((grad, i) => {
          gradMap[i] = () => grad;
        });
        return gradMap;
      };

      return this.runKernelFunc({
        forwardFunc,
        backwardsFunc,
        inputs: inputMap,
      });