Distinguish ops from kernels in tf.profile

tfjs-core/src/ops/operation.ts

@@ -16,6 +16,8 @@
  */
 import {ENGINE} from '../engine';
 
+export const OP_SCOPE_SUFFIX = '__op';
+
 /**
  * Used for wrapping functions that perform math operations on
  * Tensors. The function will be wrapped in a named scope that cleans all
@@ -38,6 +40,9 @@ export function op<T extends Function>(f: {[name: string]: T}): T {
     opName = opName.substring(0, opName.length - 1);
   }
 
+  // add an __op suffix to distinguish ops from kernels in tf.profile
+  opName = opName + OP_SCOPE_SUFFIX;
+
   // tslint:disable-next-line:no-any
   const f2 = (...args: any[]) => {
     ENGINE.startScope(opName);

tfjs-core/src/ops/operation_test.ts

@@ -22,15 +22,15 @@ describeWithFlags('operation', ALL_ENVS, () => {
     const f = () => 2;
     const opfn = op({'opName': f});
 
-    expect(opfn.name).toBe('opName');
+    expect(opfn.name).toBe('opName__op');
     expect(opfn()).toBe(2);
   });
 
   it('executes, preserves function name, strips underscore', () => {
     const f = () => 2;
     const opfn = op({'opName_': f});
 
-    expect(opfn.name).toBe('opName');
+    expect(opfn.name).toBe('opName__op');
     expect(opfn()).toBe(2);
   });
 

tfjs-core/src/ops/ops.ts

@@ -206,7 +206,7 @@ export * from './dropout';
 export * from './signal_ops_util';
 export * from './in_top_k';
 
-export {op} from './operation';
+export {op, OP_SCOPE_SUFFIX} from './operation';
 
 import {rfft} from './spectral/rfft';
 import {fft} from './spectral/fft';

tfjs/tools/custom_bundle/cli.ts

@@ -27,6 +27,8 @@ import * as fs from 'fs';
 import * as path from 'path';
 import * as mkdirp from 'mkdirp';
 
+import {OP_SCOPE_SUFFIX} from '@tensorflow/tfjs-core';
+
 import {getCustomModuleString, getCustomConverterOpsModule} from './custom_module';
 import {CustomTFJSBundleConfig, SupportedBackends} from './types';
 import {esmModuleProvider} from './esm_module_provider';
@@ -99,7 +101,17 @@ function getKernelNamesForConfig(config: CustomTFJSBundleConfig) {
   // (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.
-  return config.kernels;
+
+  // 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);
 }
 
 function getOpsForConfig(config: CustomTFJSBundleConfig) {