diff --git a/tfjs-core/src/ops/operation.ts b/tfjs-core/src/ops/operation.ts index 9aac7a251d0..09ff355cd59 100644 --- a/tfjs-core/src/ops/operation.ts +++ b/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(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); diff --git a/tfjs-core/src/ops/operation_test.ts b/tfjs-core/src/ops/operation_test.ts index a9c83a2d8c3..821b60410ee 100644 --- a/tfjs-core/src/ops/operation_test.ts +++ b/tfjs-core/src/ops/operation_test.ts @@ -22,7 +22,7 @@ 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); }); @@ -30,7 +30,7 @@ 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); }); diff --git a/tfjs-core/src/ops/ops.ts b/tfjs-core/src/ops/ops.ts index 6c1d6b81108..86c907532d9 100644 --- a/tfjs-core/src/ops/ops.ts +++ b/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'; diff --git a/tfjs/tools/custom_bundle/cli.ts b/tfjs/tools/custom_bundle/cli.ts index 1e4408639fc..fda43f10df1 100755 --- a/tfjs/tools/custom_bundle/cli.ts +++ b/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) {