@epfml/discojs 0.0.1 → 0.1.0

package/dist/client/decentralized/sec_agg.js

@@ -0,0 +1,169 @@
+"use strict";
+Object.defineProperty(exports, "__esModule", { value: true });
+exports.SecAgg = void 0;
+var tslib_1 = require("tslib");
+var immutable_1 = require("immutable");
+var msgpack_lite_1 = (0, tslib_1.__importDefault)(require("msgpack-lite"));
+var __1 = require("../..");
+var base_1 = require("./base");
+var messages = (0, tslib_1.__importStar)(require("./messages"));
+var secret_shares = (0, tslib_1.__importStar)(require("./secret_shares"));
+/**
+ * Decentralized client that utilizes secure aggregation so client updates remain private
+ */
+var SecAgg = /** @class */ (function (_super) {
+    (0, tslib_1.__extends)(SecAgg, _super);
+    function SecAgg() {
+        var _this = _super !== null && _super.apply(this, arguments) || this;
+        // list of weights received from other clients
+        _this.receivedShares = (0, immutable_1.List)();
+        // list of partial sums received by client
+        _this.receivedPartialSums = (0, immutable_1.List)();
+        // the partial sum calculated by the client
+        _this.mySum = [];
+        return _this;
+    }
+    /*
+    generates shares and sends to all ready peers adds differential privacy
+     */
+    SecAgg.prototype.sendShares = function (noisyWeights, round, trainingInformant) {
+        return (0, tslib_1.__awaiter)(this, void 0, void 0, function () {
+            var weightShares, i, weights, msg, encodedMsg;
+            var _a;
+            return (0, tslib_1.__generator)(this, function (_b) {
+                switch (_b.label) {
+                    case 0:
+                        weightShares = secret_shares.generateAllShares(noisyWeights, this.peers.length, this.maxShareValue);
+                        i = 0;
+                        _b.label = 1;
+                    case 1:
+                        if (!(i < this.peers.length)) return [3 /*break*/, 4];
+                        weights = weightShares.get(i);
+                        if (weights === undefined) {
+                            throw new Error('weight shares generated incorrectly');
+                        }
+                        _a = {
+                            type: messages.messageType.clientSharesMessageServer,
+                            peerID: this.ID
+                        };
+                        return [4 /*yield*/, __1.serialization.weights.encode(weights)];
+                    case 2:
+                        msg = (_a.weights = _b.sent(),
+                            _a.destination = this.peers[i],
+                            _a);
+                        encodedMsg = msgpack_lite_1.default.encode(msg);
+                        this.sendMessagetoPeer(encodedMsg);
+                        _b.label = 3;
+                    case 3:
+                        i++;
+                        return [3 /*break*/, 1];
+                    case 4: return [2 /*return*/];
+                }
+            });
+        });
+    };
+    /*
+  sends partial sums to connected peers so final update can be calculated
+   */
+    SecAgg.prototype.sendPartialSums = function () {
+        return (0, tslib_1.__awaiter)(this, void 0, void 0, function () {
+            var i, msg, encodedMsg;
+            var _a;
+            return (0, tslib_1.__generator)(this, function (_b) {
+                switch (_b.label) {
+                    case 0:
+                        // calculating my personal partial sum from received shares that i will share with peers
+                        this.mySum = __1.aggregation.sumWeights((0, immutable_1.List)(Array.from(this.receivedShares.values())));
+                        i = 0;
+                        _b.label = 1;
+                    case 1:
+                        if (!(i < this.peers.length)) return [3 /*break*/, 4];
+                        _a = {
+                            type: messages.messageType.clientPartialSumsMessageServer,
+                            peerID: this.ID
+                        };
+                        return [4 /*yield*/, __1.serialization.weights.encode(this.mySum)];
+                    case 2:
+                        msg = (_a.partials = _b.sent(),
+                            _a.destination = this.peers[i],
+                            _a);
+                        encodedMsg = msgpack_lite_1.default.encode(msg);
+                        this.sendMessagetoPeer(encodedMsg);
+                        _b.label = 3;
+                    case 3:
+                        i++;
+                        return [3 /*break*/, 1];
+                    case 4: return [2 /*return*/];
+                }
+            });
+        });
+    };
+    SecAgg.prototype.sendAndReceiveWeights = function (noisyWeights, round, trainingInformant) {
+        return (0, tslib_1.__awaiter)(this, void 0, void 0, function () {
+            var _this = this;
+            return (0, tslib_1.__generator)(this, function (_a) {
+                switch (_a.label) {
+                    case 0:
+                        // reset fields at beginning of each round
+                        this.receivedShares = this.receivedShares.clear();
+                        this.receivedPartialSums = this.receivedPartialSums.clear();
+                        // PHASE 1 COMMUNICATION --> send additive shares to ready peers, pause program until shares are received from all peers
+                        return [4 /*yield*/, this.sendShares(noisyWeights, round, trainingInformant)];
+                    case 1:
+                        // PHASE 1 COMMUNICATION --> send additive shares to ready peers, pause program until shares are received from all peers
+                        _a.sent();
+                        return [4 /*yield*/, this.pauseUntil(function () { return _this.receivedShares.size >= _this.peers.length; })
+                            // PHASE 2 COMMUNICATION --> send partial sums to ready peers
+                        ];
+                    case 2:
+                        _a.sent();
+                        // PHASE 2 COMMUNICATION --> send partial sums to ready peers
+                        return [4 /*yield*/, this.sendPartialSums()
+                            // after all partial sums are received, return list of partial sums to be aggregated
+                        ];
+                    case 3:
+                        // PHASE 2 COMMUNICATION --> send partial sums to ready peers
+                        _a.sent();
+                        // after all partial sums are received, return list of partial sums to be aggregated
+                        return [4 /*yield*/, this.pauseUntil(function () { return _this.receivedPartialSums.size >= _this.peers.length; })];
+                    case 4:
+                        // after all partial sums are received, return list of partial sums to be aggregated
+                        _a.sent();
+                        return [2 /*return*/, this.receivedPartialSums];
+                }
+            });
+        });
+    };
+    /*
+    checks if message contains shares from a peer
+     */
+    SecAgg.prototype.instanceOfClientSharesMessageServer = function (msg) {
+        return msg.type === messages.messageType.clientSharesMessageServer;
+    };
+    /*
+    checks if message contains partial sums from a peer
+     */
+    SecAgg.prototype.instanceOfClientPartialSumsMessageServer = function (msg) {
+        return msg.type === messages.messageType.clientPartialSumsMessageServer;
+    };
+    /*
+  handles received messages from signaling server
+   */
+    SecAgg.prototype.clientHandle = function (msg) {
+        if (this.instanceOfClientSharesMessageServer(msg)) {
+            // update received weights by one weights reception
+            var weights = __1.serialization.weights.decode(msg.weights);
+            this.receivedShares = this.receivedShares.push(weights);
+        }
+        else if (this.instanceOfClientPartialSumsMessageServer(msg)) {
+            // update received partial sums by one partial sum
+            var partials = __1.serialization.weights.decode(msg.partials);
+            this.receivedPartialSums = this.receivedPartialSums.push(partials);
+        }
+        else {
+            throw new Error('Unexpected Message Type');
+        }
+    };
+    return SecAgg;
+}(base_1.Base));
+exports.SecAgg = SecAgg;

package/dist/client/decentralized/secret_shares.d.ts

@@ -0,0 +1,5 @@
+import { List } from 'immutable';
+import { Weights } from '../..';
+export declare function lastShare(currentShares: Weights[], secret: Weights): Weights;
+export declare function generateAllShares(secret: Weights, nParticipants: number, maxShareValue: number): List<Weights>;
+export declare function generateRandomShare(secret: Weights, maxShareValue: number): Weights;

package/dist/client/decentralized/secret_shares.js

@@ -0,0 +1,58 @@
+"use strict";
+Object.defineProperty(exports, "__esModule", { value: true });
+exports.generateRandomShare = exports.generateAllShares = exports.lastShare = void 0;
+var tslib_1 = require("tslib");
+var immutable_1 = require("immutable");
+var tf = (0, tslib_1.__importStar)(require("@tensorflow/tfjs"));
+var crypto = (0, tslib_1.__importStar)(require("crypto"));
+var __1 = require("../..");
+var maxSeed = Math.pow(2, 47);
+/*
+Return Weights in the remaining share once N-1 shares have been constructed (where N is number of ready clients)
+ */
+function lastShare(currentShares, secret) {
+    if (currentShares.length === 0) {
+        throw new Error('Need at least one current share to be able to subtract secret from');
+    }
+    var currentShares2 = (0, immutable_1.List)(currentShares);
+    var last = __1.aggregation.subtractWeights((0, immutable_1.List)([secret, __1.aggregation.sumWeights(currentShares2)]));
+    return last;
+}
+exports.lastShare = lastShare;
+/*
+Generate N additive shares that aggregate to the secret weights array (where N is number of ready clients)
+ */
+function generateAllShares(secret, nParticipants, maxShareValue) {
+    var shares = [];
+    for (var i = 0; i < nParticipants - 1; i++) {
+        shares.push(generateRandomShare(secret, maxShareValue));
+    }
+    shares.push(lastShare(shares, secret));
+    var sharesFinal = (0, immutable_1.List)(shares);
+    return sharesFinal;
+}
+exports.generateAllShares = generateAllShares;
+/*
+generates one share in the same shape as the secret that is populated with values randomly chosend from
+a uniform distribution between (-maxShareValue, maxShareValue).
+ */
+function generateRandomShare(secret, maxShareValue) {
+    var e_1, _a;
+    var share = [];
+    var seed = crypto.randomInt(maxSeed);
+    try {
+        for (var secret_1 = (0, tslib_1.__values)(secret), secret_1_1 = secret_1.next(); !secret_1_1.done; secret_1_1 = secret_1.next()) {
+            var t = secret_1_1.value;
+            share.push(tf.randomUniform(t.shape, -maxShareValue, maxShareValue, 'float32', seed));
+        }
+    }
+    catch (e_1_1) { e_1 = { error: e_1_1 }; }
+    finally {
+        try {
+            if (secret_1_1 && !secret_1_1.done && (_a = secret_1.return)) _a.call(secret_1);
+        }
+        finally { if (e_1) throw e_1.error; }
+    }
+    return share;
+}
+exports.generateRandomShare = generateRandomShare;

package/dist/client/decentralized/types.d.ts

@@ -0,0 +1 @@
+export declare type PeerID = number;

package/dist/client/decentralized/types.js

@@ -0,0 +1,2 @@
+"use strict";
+Object.defineProperty(exports, "__esModule", { value: true });

package/dist/client/federated.d.ts

@@ -1,4 +1,4 @@
-import { TrainingInformant, MetadataID, Weights } from '..';
+import { informant, MetadataID, Weights } from '..';
 import { Base } from './base';
 /**
  * Class that deals with communication with the centralized server when training
@@ -24,7 +24,7 @@ export declare class Federated extends Base {
     getMetadataMap(metadataID: MetadataID): Promise<Map<string, unknown>>;
     getLatestServerRound(): Promise<number>;
     pullRoundAndFetchWeights(): Promise<Weights | undefined>;
-    pullServerStatistics(trainingInformant: TrainingInformant): Promise<void>;
-    onRoundEndCommunication(updatedWeights: Weights, staleWeights: Weights, _: number, trainingInformant: TrainingInformant): Promise<Weights>;
-    onTrainEndCommunication(_: Weights, trainingInformant: TrainingInformant): Promise<void>;
+    pullServerStatistics(trainingInformant: informant.FederatedInformant): Promise<void>;
+    onRoundEndCommunication(updatedWeights: Weights, staleWeights: Weights, _: number, trainingInformant: informant.FederatedInformant): Promise<Weights>;
+    onTrainEndCommunication(): Promise<void>;
 }

package/dist/client/federated.js

@@ -181,7 +181,7 @@ var Federated = /** @class */ (function (_super) {
                     case 0: return [4 /*yield*/, axios_1.default.get(this.urlTo('statistics'))];
                     case 1:
                         response = _a.sent();
-                        trainingInformant.updateWithServerStatistics(response.data.statistics);
+                        trainingInformant.update(response.data.statistics);
                         return [2 /*return*/];
                 }
             });
@@ -208,13 +208,10 @@ var Federated = /** @class */ (function (_super) {
             });
         });
     };
-    Federated.prototype.onTrainEndCommunication = function (_, trainingInformant) {
-        return (0, tslib_1.__awaiter)(this, void 0, void 0, function () {
-            return (0, tslib_1.__generator)(this, function (_a) {
-                trainingInformant.addMessage('Training finished.');
-                return [2 /*return*/];
-            });
-        });
+    Federated.prototype.onTrainEndCommunication = function () {
+        return (0, tslib_1.__awaiter)(this, void 0, void 0, function () { return (0, tslib_1.__generator)(this, function (_a) {
+            return [2 /*return*/];
+        }); });
     };
     return Federated;
 }(base_1.Base));

package/dist/client/index.d.ts

@@ -1,4 +1,4 @@
 export { Base } from './base';
-export { Decentralized } from './decentralized';
+export * as decentralized from './decentralized';
 export { Federated } from './federated';
 export { Local } from './local';

package/dist/client/index.js

@@ -1,10 +1,10 @@
 "use strict";
 Object.defineProperty(exports, "__esModule", { value: true });
-exports.Local = exports.Federated = exports.Decentralized = exports.Base = void 0;
+exports.Local = exports.Federated = exports.decentralized = exports.Base = void 0;
+var tslib_1 = require("tslib");
 var base_1 = require("./base");
 Object.defineProperty(exports, "Base", { enumerable: true, get: function () { return base_1.Base; } });
-var decentralized_1 = require("./decentralized");
-Object.defineProperty(exports, "Decentralized", { enumerable: true, get: function () { return decentralized_1.Decentralized; } });
+exports.decentralized = (0, tslib_1.__importStar)(require("./decentralized"));
 var federated_1 = require("./federated");
 Object.defineProperty(exports, "Federated", { enumerable: true, get: function () { return federated_1.Federated; } });
 var local_1 = require("./local");

package/dist/client/local.js

@@ -20,9 +20,11 @@ var Local = /** @class */ (function (_super) {
         }); });
     };
     Local.prototype.onRoundEndCommunication = function (_) {
-        return (0, tslib_1.__awaiter)(this, void 0, void 0, function () { return (0, tslib_1.__generator)(this, function (_a) {
-            return [2 /*return*/, _];
-        }); });
+        return (0, tslib_1.__awaiter)(this, void 0, void 0, function () {
+            return (0, tslib_1.__generator)(this, function (_a) {
+                return [2 /*return*/, _];
+            });
+        });
     };
     Local.prototype.onTrainEndCommunication = function () {
         return (0, tslib_1.__awaiter)(this, void 0, void 0, function () { return (0, tslib_1.__generator)(this, function (_a) {

package/dist/dataset/data_loader/data_loader.d.ts

@@ -3,14 +3,20 @@ import { Task } from '../../task';
 export interface DataConfig {
     features?: string[];
     labels?: string[];
+    shuffle?: boolean;
+    validationSplit?: number;
 }
 export interface Data {
     dataset: Dataset;
     size: number;
 }
+export interface DataTuple {
+    train: Data;
+    validation?: Data;
+}
 export declare abstract class DataLoader<Source> {
     protected task: Task;
     constructor(task: Task);
     abstract load(source: Source, config: DataConfig): Promise<Dataset>;
-    abstract loadAll(sources: Source[], config: DataConfig): Promise<Data>;
+    abstract loadAll(sources: Source[], config: DataConfig): Promise<DataTuple>;
 }

package/dist/dataset/data_loader/image_loader.d.ts

@@ -1,6 +1,6 @@
-import * as tf from '@tensorflow/tfjs';
+import { tf } from '../..';
 import { Dataset } from '../dataset_builder';
-import { DataLoader, DataConfig, Data } from './data_loader';
+import { DataLoader, DataConfig, DataTuple } from './data_loader';
 /**
  * TODO @s314cy:
  * Load labels and correctly match them with their respective images, with the following constraints:
@@ -10,5 +10,7 @@ import { DataLoader, DataConfig, Data } from './data_loader';
 export declare abstract class ImageLoader<Source> extends DataLoader<Source> {
     abstract readImageFrom(source: Source): Promise<tf.Tensor3D>;
     load(image: Source, config?: DataConfig): Promise<Dataset>;
-    loadAll(images: Source[], config?: DataConfig): Promise<Data>;
+    private buildDataset;
+    loadAll(images: Source[], config?: DataConfig): Promise<DataTuple>;
+    shuffle(array: number[]): void;
 }

package/dist/dataset/data_loader/image_loader.js

@@ -2,7 +2,8 @@
 Object.defineProperty(exports, "__esModule", { value: true });
 exports.ImageLoader = void 0;
 var tslib_1 = require("tslib");
-var tf = (0, tslib_1.__importStar)(require("@tensorflow/tfjs"));
+var immutable_1 = require("immutable");
+var __1 = require("../..");
 var data_loader_1 = require("./data_loader");
 /**
  * TODO @s314cy:
@@ -35,25 +36,17 @@ var ImageLoader = /** @class */ (function (_super) {
                             _a.ys = config.labels[0],
                             _a);
                         _b.label = 4;
-                    case 4: return [2 /*return*/, tf.data.array([tensorContainer])];
+                    case 4: return [2 /*return*/, __1.tf.data.array([tensorContainer])];
                 }
             });
         });
     };
-    ImageLoader.prototype.loadAll = function (images, config) {
-        var _a, _b;
+    ImageLoader.prototype.buildDataset = function (images, labels, indices, config) {
         return (0, tslib_1.__awaiter)(this, void 0, void 0, function () {
-            var labels, numberOfClasses, dataset;
+            var dataset;
             var _this = this;
-            return (0, tslib_1.__generator)(this, function (_c) {
-                if ((config === null || config === void 0 ? void 0 : config.labels) !== undefined) {
-                    numberOfClasses = (_b = (_a = this.task.trainingInformation) === null || _a === void 0 ? void 0 : _a.LABEL_LIST) === null || _b === void 0 ? void 0 : _b.length;
-                    if (numberOfClasses === undefined) {
-                        throw new Error('wanted labels but none found in task');
-                    }
-                    labels = tf.oneHot(tf.tensor1d(config.labels, 'int32'), numberOfClasses).arraySync();
-                }
-                dataset = tf.data.generator(function () {
+            return (0, tslib_1.__generator)(this, function (_a) {
+                dataset = __1.tf.data.generator(function () {
                     var withLabels = (config === null || config === void 0 ? void 0 : config.labels) !== undefined;
                     var index = 0;
                     var iterator = {
@@ -62,13 +55,13 @@ var ImageLoader = /** @class */ (function (_super) {
                             return (0, tslib_1.__generator)(this, function (_a) {
                                 switch (_a.label) {
                                     case 0:
-                                        if (index === images.length) {
+                                        if (index === indices.length) {
                                             return [2 /*return*/, { done: true }];
                                         }
-                                        return [4 /*yield*/, this.readImageFrom(images[index])];
+                                        return [4 /*yield*/, this.readImageFrom(images[indices[index]])];
                                     case 1:
                                         sample = _a.sent();
-                                        label = withLabels ? labels[index] : undefined;
+                                        label = withLabels ? labels[indices[index]] : undefined;
                                         value = withLabels ? { xs: sample, ys: label } : sample;
                                         index++;
                                         return [2 /*return*/, {
@@ -83,11 +76,64 @@ var ImageLoader = /** @class */ (function (_super) {
                 });
                 return [2 /*return*/, {
                         dataset: dataset,
-                        size: images.length
+                        size: indices.length
                     }];
             });
         });
     };
+    ImageLoader.prototype.loadAll = function (images, config) {
+        var _a, _b;
+        return (0, tslib_1.__awaiter)(this, void 0, void 0, function () {
+            var labels, indices, numberOfClasses, dataset, trainSize, trainIndices, valIndices, trainDataset, valDataset;
+            return (0, tslib_1.__generator)(this, function (_c) {
+                switch (_c.label) {
+                    case 0:
+                        labels = [];
+                        indices = (0, immutable_1.Range)(0, images.length).toArray();
+                        if ((config === null || config === void 0 ? void 0 : config.labels) !== undefined) {
+                            numberOfClasses = (_b = (_a = this.task.trainingInformation) === null || _a === void 0 ? void 0 : _a.LABEL_LIST) === null || _b === void 0 ? void 0 : _b.length;
+                            if (numberOfClasses === undefined) {
+                                throw new Error('wanted labels but none found in task');
+                            }
+                            labels = __1.tf.oneHot(__1.tf.tensor1d(config.labels, 'int32'), numberOfClasses).arraySync();
+                        }
+                        if ((config === null || config === void 0 ? void 0 : config.shuffle) === undefined || (config === null || config === void 0 ? void 0 : config.shuffle)) {
+                            this.shuffle(indices);
+                        }
+                        if (!((config === null || config === void 0 ? void 0 : config.validationSplit) === undefined)) return [3 /*break*/, 2];
+                        return [4 /*yield*/, this.buildDataset(images, labels, indices, config)];
+                    case 1:
+                        dataset = _c.sent();
+                        return [2 /*return*/, {
+                                train: dataset,
+                                validation: dataset
+                            }];
+                    case 2:
+                        trainSize = Math.floor(images.length * (1 - config.validationSplit));
+                        trainIndices = indices.slice(0, trainSize);
+                        valIndices = indices.slice(trainSize);
+                        return [4 /*yield*/, this.buildDataset(images, labels, trainIndices, config)];
+                    case 3:
+                        trainDataset = _c.sent();
+                        return [4 /*yield*/, this.buildDataset(images, labels, valIndices, config)];
+                    case 4:
+                        valDataset = _c.sent();
+                        return [2 /*return*/, {
+                                train: trainDataset,
+                                validation: valDataset
+                            }];
+                }
+            });
+        });
+    };
+    ImageLoader.prototype.shuffle = function (array) {
+        for (var i = 0; i < array.length; i++) {
+            var j = Math.floor(Math.random() * i);
+            var swap = array[i];
+            array[i] = array[j];
+            array[j] = swap;
+        }
+    };
     return ImageLoader;
 }(data_loader_1.DataLoader));
 exports.ImageLoader = ImageLoader;

package/dist/dataset/data_loader/index.d.ts

@@ -1,3 +1,3 @@
-export { Data, DataConfig, DataLoader } from './data_loader';
+export { Data, DataTuple, DataConfig, DataLoader } from './data_loader';
 export { ImageLoader } from './image_loader';
 export { TabularLoader } from './tabular_loader';

package/dist/dataset/data_loader/tabular_loader.d.ts

@@ -1,7 +1,7 @@
-import { DataLoader, DataConfig, Data } from './data_loader';
+import { DataLoader, DataConfig, DataTuple } from './data_loader';
 import { Dataset } from '../dataset_builder';
 import { Task } from '../../task';
-import * as tf from '@tensorflow/tfjs';
+import { tf } from '../..';
 export declare abstract class TabularLoader<Source> extends DataLoader<Source> {
     private readonly delimiter;
     constructor(task: Task, delimiter: string);
@@ -25,5 +25,5 @@ export declare abstract class TabularLoader<Source> extends DataLoader<Source> {
       * Creates the CSV datasets based off the given sources, then fuses them into a single CSV
       * dataset.
       */
-    loadAll(sources: Source[], config: DataConfig): Promise<Data>;
+    loadAll(sources: Source[], config: DataConfig): Promise<DataTuple>;
 }

package/dist/dataset/data_loader/tabular_loader.js

@@ -4,6 +4,8 @@ exports.TabularLoader = void 0;
 var tslib_1 = require("tslib");
 var data_loader_1 = require("./data_loader");
 var immutable_1 = require("immutable");
+// window size from which the dataset shuffling will sample
+var BUFFER_SIZE = 1000;
 var TabularLoader = /** @class */ (function (_super) {
     (0, tslib_1.__extends)(TabularLoader, _super);
     function TabularLoader(task, delimiter) {
@@ -21,7 +23,7 @@ var TabularLoader = /** @class */ (function (_super) {
      */
     TabularLoader.prototype.load = function (source, config) {
         return (0, tslib_1.__awaiter)(this, void 0, void 0, function () {
-            var columnConfigs, csvConfig;
+            var columnConfigs, csvConfig, dataset;
             return (0, tslib_1.__generator)(this, function (_a) {
                 /**
                  * Prepare the CSV config object based off the given features and labels.
@@ -40,19 +42,20 @@ var TabularLoader = /** @class */ (function (_super) {
                     configuredColumnsOnly: true,
                     delimiter: this.delimiter
                 };
-                return [2 /*return*/, this.loadTabularDatasetFrom(source, csvConfig).map(function (t) {
-                        if (typeof t === 'object' && ('xs' in t) && ('ys' in t)) {
-                            return t;
-                        }
-                        throw new Error('expected TensorContainerObject');
-                    }).map(function (t) {
-                        // TODO order may not be stable between tensor
-                        var _a = t, xs = _a.xs, ys = _a.ys;
-                        return {
-                            xs: Object.values(xs),
-                            ys: Object.values(ys)
-                        };
-                    })];
+                dataset = this.loadTabularDatasetFrom(source, csvConfig).map(function (t) {
+                    if (typeof t === 'object' && ('xs' in t) && ('ys' in t)) {
+                        return t;
+                    }
+                    throw new Error('expected TensorContainerObject');
+                }).map(function (t) {
+                    // TODO order may not be stable between tensor
+                    var _a = t, xs = _a.xs, ys = _a.ys;
+                    return {
+                        xs: Object.values(xs),
+                        ys: Object.values(ys)
+                    };
+                });
+                return [2 /*return*/, ((config === null || config === void 0 ? void 0 : config.shuffle) === undefined || (config === null || config === void 0 ? void 0 : config.shuffle)) ? dataset.shuffle(BUFFER_SIZE) : dataset];
             });
         });
     };
@@ -62,22 +65,28 @@ var TabularLoader = /** @class */ (function (_super) {
       */
     TabularLoader.prototype.loadAll = function (sources, config) {
         return (0, tslib_1.__awaiter)(this, void 0, void 0, function () {
-            var datasets, dataset;
+            var datasets, dataset, data;
             var _this = this;
             return (0, tslib_1.__generator)(this, function (_a) {
                 switch (_a.label) {
                     case 0: return [4 /*yield*/, Promise.all(sources.map(function (source) { return (0, tslib_1.__awaiter)(_this, void 0, void 0, function () { return (0, tslib_1.__generator)(this, function (_a) {
                             switch (_a.label) {
-                                case 0: return [4 /*yield*/, this.load(source, config)];
+                                case 0: return [4 /*yield*/, this.load(source, (0, tslib_1.__assign)((0, tslib_1.__assign)({}, config), { shuffle: false }))];
                                 case 1: return [2 /*return*/, _a.sent()];
                             }
                         }); }); }))];
                     case 1:
                         datasets = _a.sent();
                         dataset = (0, immutable_1.List)(datasets).reduce(function (acc, dataset) { return acc.concatenate(dataset); });
+                        data = {
+                            dataset: (config === null || config === void 0 ? void 0 : config.shuffle) ? dataset.shuffle(BUFFER_SIZE) : dataset,
+                            // dataset.size does not work for csv datasets
+                            // https://github.com/tensorflow/tfjs/issues/5845
+                            size: 0
+                        };
+                        // TODO: Implement validation split for tabular data (tricky due to streaming)
                         return [2 /*return*/, {
-                                dataset: dataset,
-                                size: dataset.size // TODO: needs to be tested
+                                train: data
                             }];
                 }
             });

package/dist/dataset/dataset_builder.d.ts

@@ -1,5 +1,5 @@
 import * as tf from '@tensorflow/tfjs';
-import { DataLoader, Data } from './data_loader/data_loader';
+import { DataConfig, DataLoader, DataTuple } from './data_loader/data_loader';
 import { Task } from '@/task';
 export declare type Dataset = tf.data.Dataset<tf.TensorContainer>;
 export declare class DatasetBuilder<Source> {
@@ -11,7 +11,8 @@ export declare class DatasetBuilder<Source> {
     constructor(dataLoader: DataLoader<Source>, task: Task);
     addFiles(sources: Source[], label?: string): void;
     clearFiles(label?: string): void;
-    build(): Promise<Data>;
+    private getLabels;
+    build(config?: DataConfig): Promise<DataTuple>;
     isBuilt(): boolean;
     size(): number;
 }