federer
Version:
Experiments in asynchronous federated learning and decentralized learning
29 lines • 1.32 kB
JavaScript
;
Object.defineProperty(exports, "__esModule", { value: true });
exports.assertNoLeakingTensors = void 0;
const tslib_1 = require("tslib");
const assert = require("assert");
const tf = tslib_1.__importStar(require("@tensorflow/tfjs-node"));
/**
* Asserts that a piece of code does not leak TensorFlow.js tensors.
*
* This is useful for debugging memory leaks. However, calls to this function
* can also be left in the code to indicate that we do not expect the total
* number of tensors in memory to grow once the code has finished executing.
*
* @param thunkName Name of the thunk; used for assertion messages
* @param thunk Function to execute
* @returns The value returned by `thunk`
*/
function assertNoLeakingTensors(thunkName, thunk) {
const numTensorsBefore = tf.memory().numTensors;
const result = thunk();
const numTensorsAfter = tf.memory().numTensors;
assert.strictEqual(numTensorsAfter, numTensorsBefore, `Found a memory leak in thunk named '${thunkName}'; ` +
`there were ${numTensorsBefore} before it was executed, ` +
`and ${numTensorsAfter} after it executed. ` +
`Expected ${numTensorsBefore} tensors in memory after execution`);
return result;
}
exports.assertNoLeakingTensors = assertNoLeakingTensors;
//# sourceMappingURL=debug.js.map