Method runKernelFunc

tfjs-core/src/engine.ts:595–744

* Internal helper method to execute a kernel Func * * Use `runKernel` to execute kernels from outside of engine.

(
      kernelParams: RegisteredKernelInvocation<I>|
      CustomGradKernelInvocation<T, I>)

   * Use `runKernel` to execute kernels from outside of engine.
   */
  private runKernelFunc<T extends Tensor|Tensor[], I extends NamedTensorMap>(
      kernelParams: RegisteredKernelInvocation<I>|
      CustomGradKernelInvocation<T, I>): T {
    let outputs: Tensor[];
    let saved: Tensor[] = [];
    const isTapeOn = this.isTapeOn();

    const startingBytecount = this.state.numBytes;
    const startingNumTensors = this.state.numTensors;

    if (this.shouldCheckForMemLeaks()) {
      this.state.numDataMovesStack.push(0);
    }

    let kernelFunc: () => Tensor[];
    if (this.backendName == null) {
      // backend has not been initialized yet (backend initialization is lazy
      // can be deferred until an op/ kernel is run).
      // The below getter has side effects that will try to initialize the
      // backend and set properties like this.backendName
      // tslint:disable-next-line: no-unused-expression
      this.backend;
    }

    let out: TensorInfo|TensorInfo[];

    const kernelOrScopeName = isRegisteredKernelInvocation(kernelParams) ?
        kernelParams.kernelName :
        this.state.activeScope != null ? this.state.activeScope.name : '';

    // Create the kernelFunc from either a registered kernel OR passed in
    // forward/backward functions (used by custom grad). In this context a
    // kernelFunc wraps a kernel implementation with some bookkeeping.

    if (isRegisteredKernelInvocation(kernelParams)) {
      const {kernelName, inputs, attrs} = kernelParams;
      if (this.backendName == null) {
        // backend has not been initialized yet (backend initialization is lazy
        // can be deferred until an op/ kernel is run).
        // The below getter has side effects that will try to initialize the
        // backend and set properties like this.backendName
        // tslint:disable-next-line: no-unused-expression
        this.backend;
      }
      const kernel = getKernel(kernelName, this.backendName);
      util.assert(
          kernel != null,
          () => `Cannot find registered kernel '${kernelName}' for backend '${
              this.backendName}'`);

      kernelFunc = () => {
        const numDataIdsBefore = this.backend.numDataIds();
        out = kernel.kernelFunc({inputs, attrs, backend: this.backend});
        const outInfos = Array.isArray(out) ? out : [out];
        if (this.shouldCheckForMemLeaks()) {
          this.checkKernelForMemLeak(kernelName, numDataIdsBefore, outInfos);
        }