federer/dist/src/common/tidy.js

Experiments in asynchronous federated learning and decentralized learning

"use strict";
Object.defineProperty(exports, "__esModule", { value: true });
exports.tidySequentialAsync = exports.tidy = void 0;
const tslib_1 = require("tslib");
const tf = tslib_1.__importStar(require("@tensorflow/tfjs-node"));
const await_lock_1 = tslib_1.__importDefault(require("await-lock"));
const dataset_1 = require("./dataset");
const weights_1 = require("./weights");
/**
 * Equivalent to `tf.tidy`, but only supporting functions returning our data
 * structures ({@link Weights}, {@link Dataset}, {@link DataSubset}, ...).
 *
 * @param fn Function to execute
 * @returns The value returned by `fn`
 */
function tidy(fn) {
    tf.engine().startScope();
    const res = fn();
    tf.engine().endScope(getTFTensorContainer(res));
    return res;
}
exports.tidy = tidy;
const tfEngineScopeLock = new await_lock_1.default();
/**
 * Equivalent to `tf.tidy`, but for functions returning our data structures
 * ({@link Weights}, {@link Dataset}, {@link DataSubset}, ...).
 *
 * Unlike `tf.tidy`, this function supports tidying up after async functions.
 * However, it does so by using a mutex lock. This means that the no two
 * promises wrapped in {@link tidySequentialAsync} can run concurrently. Note
 * that there is some slight nuance here: the promises that they create
 * internally can run concurrently, but the returned promise cannot run
 * concurrently with another wrapped promise.
 *
 * @param fn Async function to execute
 * @returns A promise of the value returned by `fn`
 */
async function tidySequentialAsync(fn) {
    // Using `tf.engine().startScope()` and `tf.engine().endScope()` in an
    // asynchronous context is generally a bad idea; see the discussion in
    // `website/docs/advanced/memory-management.md`.
    //
    // However, with a mutex lock, we can get around the problems caused by
    // interleaved microtasks. Things can happen asynchronously internally, but
    // the microtasks of two promises wrapped in this method cannot interleave,
    // so we have no interleaving of `startScope` and `endScope`; the mutex lock
    // effectively imposes a partial order on the microtask execution.
    await tfEngineScopeLock.acquireAsync();
    tf.engine().startScope();
    const res = await fn();
    tf.engine().endScope([
        res,
        getTFTensorContainer(res),
    ]);
    tfEngineScopeLock.release();
    return res;
}
exports.tidySequentialAsync = tidySequentialAsync;
/** Transforms one of our tensor containers to a `tf.TensorContainer` */
function getTFTensorContainer(result) {
    if (result === undefined) {
        return undefined;
    }
    else if (result instanceof weights_1.Weights) {
        return result.weights;
    }
    else if (result instanceof dataset_1.Dataset) {
        return [
            result.train.labels,
            result.train.items,
            result.test.labels,
            result.test.items,
        ];
    }
    else if (result instanceof dataset_1.DataSubset) {
        return [result.items, result.labels];
    }
    else if (Array.isArray(result)) {
        return result.map((item) => getTFTensorContainer(item));
    }
    else if (result instanceof Object) {
        return Object.values(result).map((item) => getTFTensorContainer(item));
    }
}
//# sourceMappingURL=tidy.js.map