From 6755bfd5232754fcd11ead7f17fb1918d4a03c0d Mon Sep 17 00:00:00 2001 From: Yannick Assogba Date: Wed, 9 Sep 2020 13:59:29 -0400 Subject: [PATCH 1/2] distinguish ops from kernels in tf.profile --- tfjs-core/src/ops/operation.ts | 4 ++++ tfjs-core/src/ops/operation_test.ts | 4 ++-- tfjs/tools/custom_bundle/cli.ts | 10 +++++++++- 3 files changed, 15 insertions(+), 3 deletions(-) diff --git a/tfjs-core/src/ops/operation.ts b/tfjs-core/src/ops/operation.ts index 9aac7a251d0..4f482426fc5 100644 --- a/tfjs-core/src/ops/operation.ts +++ b/tfjs-core/src/ops/operation.ts @@ -38,6 +38,10 @@ export function op(f: {[name: string]: T}): T { opName = opName.substring(0, opName.length - 1); } + // an an __op suffix to distinguish ops from kernels in tf.profile + + opName = opName + '__op'; + // 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/tools/custom_bundle/cli.ts b/tfjs/tools/custom_bundle/cli.ts index 1e4408639fc..72c7245b979 100755 --- a/tfjs/tools/custom_bundle/cli.ts +++ b/tfjs/tools/custom_bundle/cli.ts @@ -99,7 +99,15 @@ 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) { + return !kernelName.endsWith('__op'); + } + + return config.kernels.filter(isNotCustomOp); } function getOpsForConfig(config: CustomTFJSBundleConfig) { From 7044e5ea5ba1654074acbb6b798b35029b03e113 Mon Sep 17 00:00:00 2001 From: Yannick Assogba Date: Thu, 10 Sep 2020 12:01:05 -0400 Subject: [PATCH 2/2] code review fixes --- tfjs-core/src/ops/operation.ts | 7 ++++--- tfjs-core/src/ops/ops.ts | 2 +- tfjs/tools/custom_bundle/cli.ts | 6 +++++- 3 files changed, 10 insertions(+), 5 deletions(-) diff --git a/tfjs-core/src/ops/operation.ts b/tfjs-core/src/ops/operation.ts index 4f482426fc5..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,9 +40,8 @@ export function op(f: {[name: string]: T}): T { opName = opName.substring(0, opName.length - 1); } - // an an __op suffix to distinguish ops from kernels in tf.profile - - opName = opName + '__op'; + // 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[]) => { 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 72c7245b979..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'; @@ -104,7 +106,9 @@ function getKernelNamesForConfig(config: CustomTFJSBundleConfig) { // 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) { - return !kernelName.endsWith('__op'); + // 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);