(f: {[name: string]: T})
| 25 | * memory usage after the function is done. |
| 26 | */ |
| 27 | export function op<T extends Function>(f: {[name: string]: T}): T { |
| 28 | const keys = Object.keys(f); |
| 29 | if (keys.length !== 1) { |
| 30 | throw new Error( |
| 31 | `Please provide an object with a single key ` + |
| 32 | `(operation name) mapping to a function. Got an object with ` + |
| 33 | `${keys.length} keys.`); |
| 34 | } |
| 35 | |
| 36 | let opName = keys[0]; |
| 37 | const fn = f[opName]; |
| 38 | |
| 39 | // Strip the underscore from the end of the function name. |
| 40 | if (opName.endsWith('_')) { |
| 41 | opName = opName.substring(0, opName.length - 1); |
| 42 | } |
| 43 | |
| 44 | // add an __op suffix to distinguish ops from kernels in tf.profile |
| 45 | opName = opName + OP_SCOPE_SUFFIX; |
| 46 | |
| 47 | // tslint:disable-next-line:no-any |
| 48 | const f2 = (...args: any[]) => { |
| 49 | ENGINE.startScope(opName); |
| 50 | try { |
| 51 | const result = fn(...args); |
| 52 | if (isPromise(result)) { |
| 53 | console.error('Cannot return a Promise inside of tidy.'); |
| 54 | } |
| 55 | ENGINE.endScope(result); |
| 56 | return result; |
| 57 | } catch (ex) { |
| 58 | ENGINE.endScope(null); |
| 59 | throw ex; |
| 60 | } |
| 61 | }; |
| 62 | Object.defineProperty(f2, 'name', {value: opName, configurable: true}); |
| 63 | |
| 64 | // tslint:disable-next-line:no-any |
| 65 | return f2 as any as T; |
| 66 | } |
no outgoing calls
no test coverage detected