Function getKernelNamesForConfig

tfjs/tools/custom_module/cli.ts:118–134

(config: CustomTFJSBundleConfig)

}

function getKernelNamesForConfig(config: CustomTFJSBundleConfig) {
  // Later on this will do a union of kernels from entries, models and
  // kernels, (and kernels used by the converter itself) Currently we only
  // support directly listing kernels. remember that this also needs to handle
  // kernels used by gradients if forwardModeOnly is false.

  // Ops in core that are implemented as custom ops may appear in tf.profile
  // they will have __op as a suffix. These do not have corresponding backend
  // kernels so we need to filter them out.
  function isNotCustomOp(kernelName: string) {
    // opSuffix value is defined in tfjs-core/src/operation.ts
    // duplicating it here to avoid an export.
    return !kernelName.endsWith(OP_SCOPE_SUFFIX);
  }

  return config.kernels.filter(isNotCustomOp);
}

const customConfig = validateArgs();
const kernelsToInclude = getKernelNamesForConfig(customConfig);