diff --git a/.github/workflows/_build.yml b/.github/workflows/_build.yml
index 698d129fa..46704d40d 100644
--- a/.github/workflows/_build.yml
+++ b/.github/workflows/_build.yml
@@ -3,37 +3,6 @@ on:
   workflow_call:
 
 jobs:
-  build-lib:
-    runs-on: ubuntu-latest
-    steps:
-      - uses: actions/checkout@v7
-      - uses: ./.github/actions/setup-node
-      - run: pnpm -F discojs run build
-
-  build-lib-node:
-    needs: build-lib
-    runs-on: ubuntu-latest
-    steps:
-      - uses: actions/checkout@v7
-      - uses: ./.github/actions/setup-node
-      - run: pnpm -F discojs-node run build
-
-  build-lib-web:
-    needs: build-lib
-    runs-on: ubuntu-latest
-    steps:
-      - uses: actions/checkout@v7
-      - uses: ./.github/actions/setup-node
-      - run: pnpm -F discojs-web run build
-
-  build-server:
-    needs: build-lib-node
-    runs-on: ubuntu-latest
-    steps:
-      - uses: actions/checkout@v7
-      - uses: ./.github/actions/setup-node
-      - run: pnpm -F server run build
-
   build-server-docker:
     runs-on: ubuntu-latest
     steps:
@@ -51,29 +20,11 @@ jobs:
             timeout=$((timeout - 1))
           done
 
-  build-cli:
-    needs: build-server
+  build-topological:
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v7
       - uses: ./.github/actions/setup-node
-      # Build all subprojects to upload all artifacts at once
       - run: pnpm -r run build
       - uses: actions/upload-artifact@v7
         with: { name: all-builds, path: "*/dist" }
-
-  build-webapp:
-    needs: build-lib-web
-    runs-on: ubuntu-latest
-    steps:
-      - uses: actions/checkout@v7
-      - uses: ./.github/actions/setup-node
-      - run: pnpm -F webapp run build
-
-  build-docs-examples:
-    needs: build-server
-    runs-on: ubuntu-latest
-    steps:
-      - uses: actions/checkout@v7
-      - uses: ./.github/actions/setup-node
-      - run: pnpm -F examples run build
diff --git a/.github/workflows/_static-analysis.yml b/.github/workflows/_static-analysis.yml
index 3f59b8404..984dfb4e6 100644
--- a/.github/workflows/_static-analysis.yml
+++ b/.github/workflows/_static-analysis.yml
@@ -16,3 +16,10 @@ jobs:
       - uses: actions/checkout@v7
       - uses: ./.github/actions/setup-node
       - run: pnpm exec knip
+
+  check_cycles:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v7
+      - uses: ./.github/actions/setup-node
+      - run: pnpm run check_cycles
diff --git a/cli/src/benchmark_gpt.ts b/cli/src/benchmark_gpt.ts
index 26fb71c75..4e206317e 100644
--- a/cli/src/benchmark_gpt.ts
+++ b/cli/src/benchmark_gpt.ts
@@ -2,13 +2,13 @@ import "@tensorflow/tfjs-node";
 import { List } from "immutable";
 import { parse } from "ts-command-line-args";
 
-import type { Network, Task } from "@epfml/discojs";
+import type { Network, Task, GPTConfig } from "@epfml/discojs";
 import {
-  async_iterator,
+  gather,
   defaultTasks,
   defaultModels,
   fetchTasks,
-  models,
+  GPT,
 } from "@epfml/discojs";
 import { loadModelFromDisk, loadText } from "@epfml/discojs-node";
 
@@ -109,8 +109,8 @@ async function main(args: Required<CLIArguments>): Promise<void> {
     const epochsCount = 1;
     const iterationsPerEpoch = 10;
 
-    const config: models.GPTConfig = {
-      modelType: modelType as models.GPTConfig["modelType"],
+    const config: GPTConfig = {
+      modelType: modelType as GPTConfig["modelType"],
       maxIter: iterationsPerEpoch,
       lr: 0.0001,
       contextLength,
@@ -130,16 +130,14 @@ async function main(args: Required<CLIArguments>): Promise<void> {
       .batch(batchSize);
 
     // Init and train the model
-    const model = new models.GPT(config);
+    const model = new GPT(config);
     console.log(
       `\tmodel type ${modelType} \n\tbatch size ${batchSize} \n\tcontext length ${contextLength}`,
     );
 
     let epochTime = performance.now();
     for (let epochsCounter = 1; epochsCounter <= epochsCount; epochsCounter++) {
-      const [_, logs] = await async_iterator.gather(
-        model.train(preprocessedDataset),
-      );
+      const [_, logs] = await gather(model.train(preprocessedDataset));
       epochTime = performance.now() - epochTime;
       const msPerToken =
         epochTime /
@@ -154,7 +152,7 @@ async function main(args: Required<CLIArguments>): Promise<void> {
      */
   } else {
     const model = await loadModelFromDisk(modelPath);
-    if (!(model instanceof models.GPT)) {
+    if (!(model instanceof GPT)) {
       throw new Error("Loaded model isn't a GPT model");
     }
 
diff --git a/cli/src/cli.ts b/cli/src/cli.ts
index a2a586cc9..8ae0b7201 100644
--- a/cli/src/cli.ts
+++ b/cli/src/cli.ts
@@ -15,11 +15,7 @@ import type {
   TaskProvider,
   Network,
 } from "@epfml/discojs";
-import {
-  Disco,
-  aggregator as aggregators,
-  client as clients,
-} from "@epfml/discojs";
+import { Disco, getAggregator, getClient } from "@epfml/discojs";
 
 import { getTaskData } from "./data.js";
 import { args } from "./args.js";
@@ -35,8 +31,8 @@ async function runUser<D extends DataType, N extends Network>(
 ): Promise<List<SummaryLogs>> {
   // cast as typescript isn't good with generics
   const trainingScheme = task.trainingInformation.scheme as N;
-  const aggregator = aggregators.getAggregator(task);
-  const client = clients.getClient(trainingScheme, url, task, aggregator);
+  const aggregator = getAggregator(task);
+  const client = getClient(trainingScheme, url, task, aggregator);
   const disco = new Disco(task, client, { scheme: trainingScheme });
 
   const dir = path.join(".", `${args.testID}`);
diff --git a/cli/src/data.ts b/cli/src/data.ts
index 9079d2475..ce4ee72b3 100644
--- a/cli/src/data.ts
+++ b/cli/src/data.ts
@@ -1,6 +1,7 @@
 import path from "node:path";
-import { Dataset, processing } from "@epfml/discojs";
-import { DataFormat, DataType, Image, Task } from "@epfml/discojs";
+import type { Dataset } from "@epfml/discojs";
+import { extractColumn } from "@epfml/discojs";
+import type { DataFormat, DataType, Image, Task } from "@epfml/discojs";
 import { loadCSV, loadImage, loadImagesInDir } from "@epfml/discojs-node";
 import { Repeat } from "immutable";
 
@@ -31,10 +32,7 @@ function loadTinderDogData(split: number): Dataset<DataFormat.Raw["image"]> {
   return loadCSV(path.join(folder, "labels.csv"))
     .map(
       (row) =>
-        [
-          processing.extractColumn(row, "filename"),
-          processing.extractColumn(row, "label"),
-        ] as const,
+        [extractColumn(row, "filename"), extractColumn(row, "label")] as const,
     )
     .map(async ([filename, label]) => {
       try {
@@ -58,10 +56,7 @@ function loadData(
   return loadCSV(path.join(folder, "labels.csv"))
     .map(
       (row) =>
-        [
-          processing.extractColumn(row, "filename"),
-          processing.extractColumn(row, "label"),
-        ] as const,
+        [extractColumn(row, "filename"), extractColumn(row, "label")] as const,
     )
     .map(async ([filename, label]) => {
       try {
diff --git a/cli/src/hellaswag_gpt.ts b/cli/src/hellaswag_gpt.ts
index ba29a3221..c68862ddb 100644
--- a/cli/src/hellaswag_gpt.ts
+++ b/cli/src/hellaswag_gpt.ts
@@ -5,22 +5,25 @@ import { parse } from "ts-command-line-args";
 
 import "@tensorflow/tfjs-node";
 import path from "node:path";
-import { models, serialization, Tokenizer } from "@epfml/discojs";
+import type { HellaSwagDataset } from "@epfml/discojs";
+import {
+  GPT,
+  ONNXModel,
+  modelDecode,
+  Tokenizer,
+  evaluate_hellaswag,
+} from "@epfml/discojs";
 import { loadHellaSwag } from "@epfml/discojs-node";
 
 const __dirname = dirname(fileURLToPath(import.meta.url));
 
-async function evaluateModel(
-  model: models.GPT | models.ONNXModel,
-  numDataPoints = -1,
-) {
-  const hellaswagDataset: models.HellaSwagDataset =
-    await loadHellaSwag(numDataPoints);
+async function evaluateModel(model: GPT | ONNXModel, numDataPoints = -1) {
+  const hellaswagDataset: HellaSwagDataset = await loadHellaSwag(numDataPoints);
   const tokenizer = await Tokenizer.from_pretrained("Xenova/gpt2");
   console.log("Starting the HellaSwag benchmark...");
 
   const start = Date.now();
-  const accuracy = await models.evaluate_hellaswag(
+  const accuracy = await evaluate_hellaswag(
     model,
     tokenizer,
     hellaswagDataset,
@@ -91,15 +94,15 @@ async function main(): Promise<void> {
     { helpArg: "help" },
   );
 
-  let model: models.GPT | models.ONNXModel | undefined;
+  let model: GPT | ONNXModel | undefined;
   switch (args.model) {
     case "onnx":
       console.log("Using ONNX pretrained model Xenova/gpt2");
-      model = await models.ONNXModel.init_pretrained("Xenova/gpt2");
+      model = await ONNXModel.init_pretrained("Xenova/gpt2");
       break;
     case "gpt-tfjs-random":
       console.log("Using GPT-TFJS with random initialization");
-      model = new models.GPT({ seed: 42 });
+      model = new GPT({ seed: 42 });
       break;
     case "gpt-tfjs-pretrained":
       console.log("Using GPT-TFJS with pretrained weights");
@@ -109,7 +112,7 @@ async function main(): Promise<void> {
         );
       }
       const encodedModel = await fsPromise.readFile(args.pretrainedModelPath);
-      model = (await serialization.model.decode(encodedModel)) as models.GPT;
+      model = (await modelDecode(encodedModel)) as GPT;
       break;
   }
   await evaluateModel(model, args.numDataPoints);
diff --git a/cli/src/train_gpt.ts b/cli/src/train_gpt.ts
index f6445f4ff..f478db14f 100644
--- a/cli/src/train_gpt.ts
+++ b/cli/src/train_gpt.ts
@@ -1,12 +1,13 @@
 import "@tensorflow/tfjs-node";
-import { models, Dataset, Tokenizer } from "@epfml/discojs";
+import type { GPTConfig } from "@epfml/discojs";
+import { GPT, Dataset, Tokenizer } from "@epfml/discojs";
 import { List } from "immutable";
 
 async function main(): Promise<void> {
   const data = "Lorem ipsum dolor sit amet, consectetur adipis";
   const seed = 42;
 
-  const config: models.GPTConfig = {
+  const config: GPTConfig = {
     modelType: "gpt-nano",
     lr: 0.01,
     maxIter: 50,
@@ -26,7 +27,7 @@ async function main(): Promise<void> {
     .repeat()
     .batch(8);
 
-  const model = new models.GPT(config);
+  const model = new GPT(config);
   for await (const logs of model.train(tokenDataset, undefined)) {
     console.log(logs);
   }
diff --git a/cli/src/user_log.ts b/cli/src/user_log.ts
index 2b7fc7bca..cdb33f015 100644
--- a/cli/src/user_log.ts
+++ b/cli/src/user_log.ts
@@ -1,4 +1,5 @@
-import { args, BenchmarkArguments } from "./args.js";
+import type { BenchmarkArguments } from "./args.js";
+import { args } from "./args.js";
 import type { SummaryLogs, DataType, Network, Task } from "@epfml/discojs";
 
 type SerializableArguments = Omit<BenchmarkArguments, "provider" | "host"> & {
diff --git a/discojs-node/package.json b/discojs-node/package.json
index 9d941a2c6..acfdc14a4 100644
--- a/discojs-node/package.json
+++ b/discojs-node/package.json
@@ -10,7 +10,7 @@
   "types": "dist/index.d.ts",
   "scripts": {
     "watch": "nodemon --ext ts --ignore dist --watch ../discojs/dist --watch . --exec pnpm run",
-    "build": "tsc --build tsconfig.lib.json",
+    "build": "tsc --build",
     "test": "cd .. && vitest --run --project=discojs-node"
   },
   "repository": {
diff --git a/discojs-node/src/hellaswag.ts b/discojs-node/src/hellaswag.ts
index 77aa23998..96577372e 100644
--- a/discojs-node/src/hellaswag.ts
+++ b/discojs-node/src/hellaswag.ts
@@ -2,7 +2,8 @@ import path from "node:path";
 import fetch from "node-fetch";
 import fs from "node:fs/promises";
 
-import { models } from "@epfml/discojs";
+import type { HellaSwagExample, HellaSwagDataset } from "@epfml/discojs";
+import { HELLASWAG_URL } from "@epfml/discojs";
 
 import { dirname } from "path";
 import { fileURLToPath } from "url";
@@ -17,7 +18,7 @@ const hellaswag_filepath = path.join(DATASET_DIR, "hellaswag_val.jsonl");
  * @param limit - Maximum number of examples to load (-1 means all)
  * @returns A HellaSwagDataset containing the examples.
  */
-export async function load(limit = -1): Promise<models.HellaSwagDataset> {
+export async function load(limit = -1): Promise<HellaSwagDataset> {
   let text: string;
   try {
     // Reads the file if it exists locally
@@ -25,10 +26,10 @@ export async function load(limit = -1): Promise<models.HellaSwagDataset> {
   } catch {
     console.log("Downloading the Hellaswag benchmark");
     // Otherwise fetch it
-    const response = await fetch(models.HELLASWAG_URL);
+    const response = await fetch(HELLASWAG_URL);
     if (!response.ok) {
       throw new Error(
-        `Failed to fetch dataset from ${models.HELLASWAG_URL}: ${response.statusText}`,
+        `Failed to fetch dataset from ${HELLASWAG_URL}: ${response.statusText}`,
       );
     }
 
@@ -39,14 +40,14 @@ export async function load(limit = -1): Promise<models.HellaSwagDataset> {
 
   const lines = text.split("\n");
 
-  const dataset: models.HellaSwagDataset = [];
+  const dataset: HellaSwagDataset = [];
   let count = 0;
   for (const line of lines) {
     if (line.trim().length === 0) continue;
     if (limit !== -1 && count >= limit) break;
 
     try {
-      const data = JSON.parse(line.trim()) as models.HellaSwagExample;
+      const data = JSON.parse(line.trim()) as HellaSwagExample;
       dataset.push(data);
       count++;
     } catch (e) {
diff --git a/discojs-node/src/loaders/text.ts b/discojs-node/src/loaders/text.ts
index cf5d22a48..c8c625e28 100644
--- a/discojs-node/src/loaders/text.ts
+++ b/discojs-node/src/loaders/text.ts
@@ -1,6 +1,7 @@
 import createDebug from "debug";
 import { createReadStream } from "node:fs";
-import { Dataset, Text } from "@epfml/discojs";
+import type { Text } from "@epfml/discojs";
+import { Dataset } from "@epfml/discojs";
 
 const debug = createDebug("discojs-node:loaders:text");
 
diff --git a/discojs-node/src/model_loader.ts b/discojs-node/src/model_loader.ts
index 1da84def0..879866dd5 100644
--- a/discojs-node/src/model_loader.ts
+++ b/discojs-node/src/model_loader.ts
@@ -1,14 +1,14 @@
 import fs from "node:fs/promises";
 
-import type { models, DataType } from "@epfml/discojs";
-import { serialization } from "@epfml/discojs";
+import type { Model, DataType } from "@epfml/discojs";
+import { modelEncode, modelDecode } from "@epfml/discojs";
 
 export async function saveModelToDisk(
-  model: models.Model<DataType>,
+  model: Model<DataType>,
   modelFolder: string,
   modelFileName: string,
 ): Promise<void> {
-  const encoded = await serialization.model.encode(model);
+  const encoded = await modelEncode(model);
 
   await fs.mkdir(modelFolder, { recursive: true });
   await fs.writeFile(`${modelFolder}/${modelFileName}`, encoded);
@@ -16,8 +16,8 @@ export async function saveModelToDisk(
 
 export async function loadModelFromDisk(
   modelPath: string,
-): Promise<models.Model<DataType>> {
+): Promise<Model<DataType>> {
   const content = await fs.readFile(modelPath);
 
-  return await serialization.model.decode(content);
+  return await modelDecode(content);
 }
diff --git a/discojs-node/tsconfig.json b/discojs-node/tsconfig.json
index 3a4772a71..7b7cf6dca 100644
--- a/discojs-node/tsconfig.json
+++ b/discojs-node/tsconfig.json
@@ -1,7 +1,4 @@
 {
-  "compilerOptions": {
-    "composite": true
-  },
   "files": [],
   "references": [
     {
diff --git a/discojs-node/tsconfig.lib.json b/discojs-node/tsconfig.lib.json
index c9d66d4a8..09fdf2e47 100644
--- a/discojs-node/tsconfig.lib.json
+++ b/discojs-node/tsconfig.lib.json
@@ -7,9 +7,7 @@
   ],
   "compilerOptions": {
     "rootDir": "./src",
-    "outDir": "dist",
-    "composite": true
+    "outDir": "dist"
   },
-  "include": ["src"],
-  "exclude": ["**/*.spec.ts"]
+  "include": ["src"]
 }
diff --git a/discojs-node/tsconfig.vitest.json b/discojs-node/tsconfig.vitest.json
index 48ca20af3..f69d69ef1 100644
--- a/discojs-node/tsconfig.vitest.json
+++ b/discojs-node/tsconfig.vitest.json
@@ -2,9 +2,11 @@
   "extends": "../tsconfig.base.json",
   "references": [
     {
-      "path": "../discojs/tsconfig.vitest.json"
+      "path": "../discojs/tsconfig.lib.json"
     }
   ],
-  "compilerOptions": { "noEmit": true },
+  "compilerOptions": {
+    "noEmit": true
+  },
   "include": ["src"]
 }
diff --git a/discojs-web/package.json b/discojs-web/package.json
index c75b07f0d..9bb48dafe 100644
--- a/discojs-web/package.json
+++ b/discojs-web/package.json
@@ -10,7 +10,7 @@
   "types": "dist/index.d.ts",
   "scripts": {
     "watch": "nodemon --ext ts --ignore dist --watch ../discojs/dist --watch . --exec pnpm run",
-    "build": "tsc --build tsconfig.lib.json",
+    "build": "tsc --build",
     "test": "cd .. && vitest --run --project=discojs-web"
   },
   "repository": {
diff --git a/discojs-web/src/hellaswag.spec.ts b/discojs-web/src/hellaswag.spec.ts
index 76b5dbcef..0c1d50c25 100644
--- a/discojs-web/src/hellaswag.spec.ts
+++ b/discojs-web/src/hellaswag.spec.ts
@@ -1,10 +1,10 @@
 import { describe, it, expect } from "vitest";
 import { load as loadHellaSwag } from "./hellaswag.js";
-import { models } from "@epfml/discojs";
+import type { HellaSwagDataset } from "@epfml/discojs";
 
 describe("hellaswag parser", () => {
   it("loads the whole hellaswag dataset", async () => {
-    const dataset: models.HellaSwagDataset = await loadHellaSwag(2);
+    const dataset: HellaSwagDataset = await loadHellaSwag(2);
 
     // basic assertions
     expect(dataset).to.be.an("array");
diff --git a/discojs-web/src/hellaswag.ts b/discojs-web/src/hellaswag.ts
index da96aac95..47e6bfe66 100644
--- a/discojs-web/src/hellaswag.ts
+++ b/discojs-web/src/hellaswag.ts
@@ -1,4 +1,5 @@
-import { models } from "@epfml/discojs";
+import type { HellaSwagExample, HellaSwagDataset } from "@epfml/discojs";
+import { HELLASWAG_URL } from "@epfml/discojs";
 
 /**
  * Loads the HellaSwag dataset from the remote URL in the browser
@@ -6,25 +7,25 @@ import { models } from "@epfml/discojs";
  * @param limit - Maximum number of examples to load (-1 means all)
  * @returns A HellaSwagDataset containing the examples
  */
-export async function load(limit = -1): Promise<models.HellaSwagDataset> {
-  const response = await fetch(models.HELLASWAG_URL);
+export async function load(limit = -1): Promise<HellaSwagDataset> {
+  const response = await fetch(HELLASWAG_URL);
   if (!response.ok) {
     throw new Error(
-      `Failed to fetch dataset from ${models.HELLASWAG_URL}: ${response.statusText}`,
+      `Failed to fetch dataset from ${HELLASWAG_URL}: ${response.statusText}`,
     );
   }
 
   const text = await response.text();
   const lines = text.split("\n");
 
-  const dataset: models.HellaSwagDataset = [];
+  const dataset: HellaSwagDataset = [];
   let count = 0;
   for (const line of lines) {
     if (line.trim().length === 0) continue;
     if (limit !== -1 && count >= limit) break;
 
     try {
-      const data = JSON.parse(line.trim()) as models.HellaSwagExample;
+      const data = JSON.parse(line.trim()) as HellaSwagExample;
       dataset.push(data);
       count++;
     } catch (e) {
diff --git a/discojs-web/src/loaders/text.ts b/discojs-web/src/loaders/text.ts
index d9ad0f89e..0c84f0150 100644
--- a/discojs-web/src/loaders/text.ts
+++ b/discojs-web/src/loaders/text.ts
@@ -1,4 +1,5 @@
-import { Dataset, Text } from "@epfml/discojs";
+import type { Text } from "@epfml/discojs";
+import { Dataset } from "@epfml/discojs";
 
 export function load(file: Blob): Dataset<Text> {
   return new Dataset(async function* () {
diff --git a/discojs-web/tsconfig.json b/discojs-web/tsconfig.json
index 3a4772a71..7b7cf6dca 100644
--- a/discojs-web/tsconfig.json
+++ b/discojs-web/tsconfig.json
@@ -1,7 +1,4 @@
 {
-  "compilerOptions": {
-    "composite": true
-  },
   "files": [],
   "references": [
     {
diff --git a/discojs-web/tsconfig.lib.json b/discojs-web/tsconfig.lib.json
index b08d8536a..09fdf2e47 100644
--- a/discojs-web/tsconfig.lib.json
+++ b/discojs-web/tsconfig.lib.json
@@ -6,11 +6,8 @@
     }
   ],
   "compilerOptions": {
-    "lib": ["DOM"],
     "rootDir": "./src",
-    "outDir": "dist",
-    "composite": true
+    "outDir": "dist"
   },
-  "include": ["src"],
-  "exclude": ["**/*.spec.ts"]
+  "include": ["src"]
 }
diff --git a/discojs-web/tsconfig.vitest.json b/discojs-web/tsconfig.vitest.json
index bfc451458..00ae1aad2 100644
--- a/discojs-web/tsconfig.vitest.json
+++ b/discojs-web/tsconfig.vitest.json
@@ -2,13 +2,12 @@
   "extends": "../tsconfig.base.json",
   "references": [
     {
-      "path": "../discojs/tsconfig.vitest.json"
+      "path": "../discojs/tsconfig.lib.json"
     }
   ],
   "compilerOptions": {
     "lib": ["DOM"],
-    "noEmit": true,
-    "composite": true
+    "noEmit": true
   },
   "include": ["src"]
 }
diff --git a/discojs/package.json b/discojs/package.json
index 93c15c651..53839572b 100644
--- a/discojs/package.json
+++ b/discojs/package.json
@@ -12,7 +12,7 @@
   "types": "dist/index.d.ts",
   "scripts": {
     "watch": "nodemon --ext ts --ignore dist --exec pnpm run",
-    "build": "tsc --build tsconfig.lib.json",
+    "build": "tsc --build",
     "test": "cd .. && vitest --run --project=discojs"
   },
   "repository": {
@@ -23,30 +23,70 @@
     "url": "https://github.com/epfml/disco/issues"
   },
   "imports": {
+    "#root/*": {
+      "@disco/source": "./src/*.ts",
+      "types": "./dist/*.d.ts",
+      "default": "./dist/*.js"
+    },
+    "#aggregator/*": {
+      "@disco/source": "./src/aggregator/*.ts",
+      "types": "./dist/aggregator/*.d.ts",
+      "default": "./dist/aggregator/*.js"
+    },
+    "#client/*": {
+      "@disco/source": "./src/client/*.ts",
+      "types": "./dist/client/*.d.ts",
+      "default": "./dist/client/*.js"
+    },
+    "#dataset/*": {
+      "@disco/source": "./src/dataset/*.ts",
+      "types": "./dist/dataset/*.d.ts",
+      "default": "./dist/dataset/*.js"
+    },
+    "#types/*": {
+      "@disco/source": "./src/types/*.ts",
+      "types": "./dist/types/*.d.ts",
+      "default": "./dist/types/*.js"
+    },
+    "#logging/*": {
+      "@disco/source": "./src/logging/*.ts",
+      "types": "./dist/logging/*.d.ts",
+      "default": "./dist/logging/*.js"
+    },
     "#models/*": {
       "@disco/source": "./src/models/*.ts",
       "types": "./dist/models/*.d.ts",
       "default": "./dist/models/*.js"
     },
+    "#processing/*": {
+      "@disco/source": "./src/processing/*.ts",
+      "types": "./dist/processing/*.d.ts",
+      "default": "./dist/processing/*.js"
+    },
+    "#serialization/*": {
+      "@disco/source": "./src/serialization/*.ts",
+      "types": "./dist/serialization/*.d.ts",
+      "default": "./dist/serialization/*.js"
+    },
     "#task/*": {
       "@disco/source": "./src/task/*.ts",
       "types": "./dist/task/*.d.ts",
       "default": "./dist/task/*.js"
     },
+    "#training/*": {
+      "@disco/source": "./src/training/*.ts",
+      "types": "./dist/training/*.d.ts",
+      "default": "./dist/training/*.js"
+    },
+    "#utils/*": {
+      "@disco/source": "./src/utils/*.ts",
+      "types": "./dist/utils/*.d.ts",
+      "default": "./dist/utils/*.js"
+    },
     "#weights/*": {
       "@disco/source": "./src/weights/*.ts",
       "types": "./dist/weights/*.d.ts",
       "default": "./dist/weights/*.js"
-    },
-    "#types/*": {
-      "@disco/source": "./src/types/*.ts",
-      "types": "./dist/types/*.d.ts",
-      "default": "./dist/types/*.js"
-    },
-    "#dataset/*": {
-      "@disco/source": "./src/dataset/*.ts",
-      "types": "./dist/dataset/*.d.ts",
-      "default": "./dist/dataset/*.js"
     }
   },
   "homepage": "https://github.com/epfml/disco#readme",
diff --git a/discojs/src/aggregator.spec.ts b/discojs/src/aggregator.spec.ts
index 40f5806e0..fb2d960e9 100644
--- a/discojs/src/aggregator.spec.ts
+++ b/discojs/src/aggregator.spec.ts
@@ -4,9 +4,9 @@ import {
   type Aggregator,
   MeanAggregator,
   SecureAggregator,
-} from "./aggregator/index.js";
-import type { NodeID } from "./client/types.js";
-import { WeightsContainer } from "./index.js";
+} from "#aggregator/index";
+import type { NodeID } from "#client/index";
+import { WeightsContainer } from "#weights/index";
 
 const AGGREGATORS: Set<[name: string, new () => Aggregator]> = Set.of<
   new () => Aggregator
diff --git a/discojs/src/aggregator/aggregator.ts b/discojs/src/aggregator/aggregator.ts
index 7f437a27d..a48f88205 100644
--- a/discojs/src/aggregator/aggregator.ts
+++ b/discojs/src/aggregator/aggregator.ts
@@ -1,9 +1,10 @@
 import createDebug from "debug";
 import { Map, Set } from "immutable";
 
-import type { client, WeightsContainer } from "../index.js";
+import type { WeightsContainer } from "#weights/index";
+import type { NodeID } from "#client/types";
 
-import { EventEmitter } from "../utils/event_emitter.js";
+import { EventEmitter } from "#utils/event_emitter";
 
 const debug = createDebug("discojs:aggregator");
 
@@ -27,14 +28,14 @@ export abstract class Aggregator extends EventEmitter<{
    * Contains the ids of all active nodes, i.e. members of the aggregation group at
    * a given round. It is a subset of all the nodes available in the network.
    */
-  protected _nodes: Set<client.NodeID>;
+  protected _nodes: Set<NodeID>;
   /**
    * Contains the contributions received from active nodes, accessible by node id.
    * It defines the effective aggregation group, which is possibly a subset
    * of all active nodes, depending on the aggregation scheme.
    */
   // communication round -> NodeID -> WeightsContainer
-  protected contributions: Map<number, Map<client.NodeID, WeightsContainer>>;
+  protected contributions: Map<number, Map<NodeID, WeightsContainer>>;
 
   /**
    * The current aggregation round, used for assessing whether a node contribution is recent enough
@@ -88,7 +89,7 @@ export abstract class Aggregator extends EventEmitter<{
    * @param contribution The node's contribution
    */
   add(
-    nodeId: client.NodeID,
+    nodeId: NodeID,
     contribution: WeightsContainer,
     aggregationRound: number,
     communicationRound?: number,
@@ -125,7 +126,7 @@ export abstract class Aggregator extends EventEmitter<{
   // Abstract method to be implemented by subclasses
   // Handles logging and adding the contribution to the list of the current round's contributions
   protected abstract _add(
-    nodeId: client.NodeID,
+    nodeId: NodeID,
     contribution: WeightsContainer,
     communicationRound?: number,
   ): void;
@@ -137,7 +138,7 @@ export abstract class Aggregator extends EventEmitter<{
    * @param nodeId the node id of the contribution to be added
    * @param round the aggregation round of the contribution to be added
    */
-  isValidContribution(nodeId: client.NodeID, round: number): boolean {
+  isValidContribution(nodeId: NodeID, round: number): boolean {
     if (!this.nodes.has(nodeId)) {
       debug("Contribution rejected because node id is not registered");
       return false;
@@ -172,7 +173,7 @@ export abstract class Aggregator extends EventEmitter<{
    * @param step The aggregation step
    * @param from The node which triggered the logging message
    */
-  log(step: AggregationStep, from?: client.NodeID): void {
+  log(step: AggregationStep, from?: NodeID): void {
     switch (step) {
       case AggregationStep.ADD:
         debug(
@@ -206,7 +207,7 @@ export abstract class Aggregator extends EventEmitter<{
    * @param nodeId The node to be added
    * @returns True is the node wasn't already in the list of nodes, False if already included
    */
-  registerNode(nodeId: client.NodeID): boolean {
+  registerNode(nodeId: NodeID): boolean {
     if (!this.nodes.has(nodeId)) {
       this._nodes = this._nodes.add(nodeId);
       return true;
@@ -218,7 +219,7 @@ export abstract class Aggregator extends EventEmitter<{
    * Remove a node's id from the set of active nodes.
    * @param nodeId The node to be removed
    */
-  removeNode(nodeId: client.NodeID): void {
+  removeNode(nodeId: NodeID): void {
     this._nodes = this._nodes.delete(nodeId);
   }
 
@@ -228,7 +229,7 @@ export abstract class Aggregator extends EventEmitter<{
    * during this aggregation round.
    * @param nodeIds The new set of nodes
    */
-  setNodes(nodeIds: Set<client.NodeID>): void {
+  setNodes(nodeIds: Set<NodeID>): void {
     this._nodes = nodeIds;
   }
 
@@ -247,16 +248,14 @@ export abstract class Aggregator extends EventEmitter<{
    * Constructs the payloads sent to other nodes as contribution.
    * @param base Object from which the payload is computed
    */
-  abstract makePayloads(
-    base: WeightsContainer,
-  ): Map<client.NodeID, WeightsContainer>;
+  abstract makePayloads(base: WeightsContainer): Map<NodeID, WeightsContainer>;
 
   abstract isFull(): boolean;
 
   /**
    * The set of node ids, representing our neighbors within the network.
    */
-  get nodes(): Set<client.NodeID> {
+  get nodes(): Set<NodeID> {
     return this._nodes;
   }
 
diff --git a/discojs/src/aggregator/byzantine.spec.ts b/discojs/src/aggregator/byzantine.spec.ts
index 3d6e1e6f6..404ceae39 100644
--- a/discojs/src/aggregator/byzantine.spec.ts
+++ b/discojs/src/aggregator/byzantine.spec.ts
@@ -2,8 +2,8 @@ import { Set } from "immutable";
 import { describe, expect, it } from "vitest";
 import fc from "fast-check";
 
-import { WeightsContainer } from "../index.js";
-import { ByzantineRobustAggregator } from "./byzantine.js";
+import { WeightsContainer } from "#weights/index";
+import { ByzantineRobustAggregator } from "#aggregator/byzantine";
 
 // Helper to convert WeightsContainer → number[][] for easy assertions
 async function WSIntoArrays(ws: WeightsContainer): Promise<number[][]> {
diff --git a/discojs/src/aggregator/byzantine.ts b/discojs/src/aggregator/byzantine.ts
index 476ef90cc..63b882669 100644
--- a/discojs/src/aggregator/byzantine.ts
+++ b/discojs/src/aggregator/byzantine.ts
@@ -1,9 +1,13 @@
 import { Map } from "immutable";
 import * as tf from "@tensorflow/tfjs";
-import { AggregationStep } from "./aggregator.js";
-import { MultiRoundAggregator, ThresholdType } from "./multiround.js";
-import { WeightsContainer, client } from "../index.js";
-import { aggregation } from "../index.js";
+
+import type { WeightsContainer } from "#weights/index";
+import type { NodeID } from "#client/types";
+import { avg } from "#weights/index";
+
+import { AggregationStep } from "#aggregator/aggregator";
+import type { ThresholdType } from "#aggregator/multiround";
+import { MultiRoundAggregator } from "#aggregator/multiround";
 
 /**
  * Byzantine-robust aggregator using Centered Clipping (CC), based on the
@@ -41,7 +45,7 @@ export class ByzantineRobustAggregator extends MultiRoundAggregator {
   private readonly clippingRadius: number;
   private readonly maxIterations: number;
   private readonly beta: number;
-  private historyMomentums: Map<client.NodeID, WeightsContainer> = Map();
+  private historyMomentums: Map<NodeID, WeightsContainer> = Map();
   private prevAggregate: WeightsContainer | null = null;
 
   /** 
@@ -86,7 +90,7 @@ export class ByzantineRobustAggregator extends MultiRoundAggregator {
     this.beta = beta;
   }
 
-  override _add(nodeId: client.NodeID, contribution: WeightsContainer): void {
+  override _add(nodeId: NodeID, contribution: WeightsContainer): void {
     this.log(
       this.contributions.hasIn([0, nodeId])
         ? AggregationStep.UPDATE
@@ -114,7 +118,7 @@ export class ByzantineRobustAggregator extends MultiRoundAggregator {
 
     // If clipping radius is infinite, fall back to simple mean
     if (!isFinite(this.clippingRadius)) {
-      return aggregation.avg(currentContributions.values());
+      return avg(currentContributions.values());
     }
 
     // Step 1: Initialize v using previous aggregate or mean of contributions
@@ -122,7 +126,7 @@ export class ByzantineRobustAggregator extends MultiRoundAggregator {
     if (this.prevAggregate) {
       v = this.prevAggregate.map((t) => tf.clone(t)); // Clone to avoid in-place modifications
     } else {
-      v = aggregation.avg(currentContributions.values());
+      v = avg(currentContributions.values());
     }
 
     const eps = tf.scalar(1e-12);
@@ -151,7 +155,7 @@ export class ByzantineRobustAggregator extends MultiRoundAggregator {
         },
       );
 
-      const avgClip = aggregation.avg(clippedDiffs);
+      const avgClip = avg(clippedDiffs);
       const newV = v.add(avgClip);
 
       clippedDiffs.forEach((d) => d.dispose());
@@ -169,7 +173,7 @@ export class ByzantineRobustAggregator extends MultiRoundAggregator {
 
   override makePayloads(
     weights: WeightsContainer,
-  ): Map<client.NodeID, WeightsContainer> {
+  ): Map<NodeID, WeightsContainer> {
     // Communicate our local weights to every other node, be it a peer or a server
     return this.nodes.toMap().map(() => weights);
   }
diff --git a/discojs/src/aggregator/byzantine_vs_percentile.spec.ts b/discojs/src/aggregator/byzantine_vs_percentile.spec.ts
index 201e425a9..929864e22 100644
--- a/discojs/src/aggregator/byzantine_vs_percentile.spec.ts
+++ b/discojs/src/aggregator/byzantine_vs_percentile.spec.ts
@@ -1,9 +1,9 @@
 import { Set } from "immutable";
 import { describe, expect, it } from "vitest";
 
-import { WeightsContainer } from "../index.js";
-import { ByzantineRobustAggregator } from "./byzantine.js";
-import { PercentileClippingAggregator } from "./percentile_clipping.js";
+import { WeightsContainer } from "#weights/index";
+import { ByzantineRobustAggregator } from "#aggregator/byzantine";
+import { PercentileClippingAggregator } from "#aggregator/percentile_clipping";
 
 // Helper to convert WeightsContainer → number[][] for easy assertions
 async function WSIntoArrays(ws: WeightsContainer): Promise<number[][]> {
diff --git a/discojs/src/aggregator/get.ts b/discojs/src/aggregator/get.ts
index ee53e1721..af5919263 100644
--- a/discojs/src/aggregator/get.ts
+++ b/discojs/src/aggregator/get.ts
@@ -1,6 +1,9 @@
-import type { DataType, Network, Task } from "../index.js";
-import { aggregator } from "../index.js";
-import { ByzantineRobustAggregator } from "./byzantine.js";
+import type { DataType, Network } from "#types/index";
+import type { Task } from "#task/index";
+import type { Aggregator } from "#aggregator/aggregator";
+import { MeanAggregator } from "#aggregator/mean";
+import { SecureAggregator } from "#aggregator/secure";
+import { ByzantineRobustAggregator } from "#aggregator/byzantine";
 
 type AggregatorOptions = Partial<{
   scheme: Task<DataType, Network>["trainingInformation"]["scheme"]; // if undefined, fallback on task.trainingInformation.scheme
@@ -30,7 +33,7 @@ type AggregatorOptions = Partial<{
 export function getAggregator(
   task: Task<DataType, Network>,
   options: AggregatorOptions = {},
-): aggregator.Aggregator {
+): Aggregator {
   const scheme = options.scheme ?? task.trainingInformation.scheme;
 
   // If options are not specified, we default to expecting a contribution from all peers, so we set the threshold to 100%
@@ -65,7 +68,7 @@ export function getAggregator(
       );
     }
     case "mean":
-      return new aggregator.MeanAggregator(
+      return new MeanAggregator(
         networkOptions.roundCutOff,
         networkOptions.threshold,
         networkOptions.thresholdType,
@@ -76,8 +79,6 @@ export function getAggregator(
           "secure aggregation is currently supported for decentralized only",
         );
       }
-      return new aggregator.SecureAggregator(
-        task.trainingInformation.maxShareValue,
-      );
+      return new SecureAggregator(task.trainingInformation.maxShareValue);
   }
 }
diff --git a/discojs/src/aggregator/index.ts b/discojs/src/aggregator/index.ts
index 695c60d62..75be361d5 100644
--- a/discojs/src/aggregator/index.ts
+++ b/discojs/src/aggregator/index.ts
@@ -1,6 +1,5 @@
-export { Aggregator, AggregationStep } from "./aggregator.js";
+export { Aggregator } from "./aggregator.js";
 export { MeanAggregator } from "./mean.js";
 export { SecureAggregator } from "./secure.js";
-export { ByzantineRobustAggregator } from "./byzantine.js";
 
 export { getAggregator } from "./get.js";
diff --git a/discojs/src/aggregator/mean.spec.ts b/discojs/src/aggregator/mean.spec.ts
index c9eab512f..6184171c6 100644
--- a/discojs/src/aggregator/mean.spec.ts
+++ b/discojs/src/aggregator/mean.spec.ts
@@ -1,7 +1,7 @@
 import { Set } from "immutable";
 import { describe, expect, it } from "vitest";
-import { WeightsContainer } from "../index.js";
-import { MeanAggregator } from "./mean.js";
+import { WeightsContainer } from "#weights/index";
+import { MeanAggregator } from "#aggregator/mean";
 
 async function WSIntoArrays(ws: WeightsContainer): Promise<number[][]> {
   return (await Promise.all(ws.weights.map(async (w) => await w.data()))).map(
diff --git a/discojs/src/aggregator/mean.ts b/discojs/src/aggregator/mean.ts
index a4abba6c4..0f1faccfd 100644
--- a/discojs/src/aggregator/mean.ts
+++ b/discojs/src/aggregator/mean.ts
@@ -1,8 +1,12 @@
 import type { Map } from "immutable";
-import { AggregationStep } from "./aggregator.js";
-import { MultiRoundAggregator, ThresholdType } from "./multiround.js";
-import type { WeightsContainer, client } from "../index.js";
-import { aggregation } from "../index.js";
+
+import type { WeightsContainer } from "#weights/index";
+import type { NodeID } from "#client/types";
+import { avg } from "#weights/index";
+
+import { AggregationStep } from "#aggregator/aggregator";
+import type { ThresholdType } from "#aggregator/multiround";
+import { MultiRoundAggregator } from "#aggregator/multiround";
 
 /**
  * Mean aggregator whose aggregation step consists in computing the mean of the received weights.
@@ -18,7 +22,7 @@ export class MeanAggregator extends MultiRoundAggregator {
     super(roundCutoff, threshold, thresholdType);
   }
 
-  override _add(nodeId: client.NodeID, contribution: WeightsContainer): void {
+  override _add(nodeId: NodeID, contribution: WeightsContainer): void {
     this.log(
       this.contributions.hasIn([0, nodeId])
         ? AggregationStep.UPDATE
@@ -35,13 +39,13 @@ export class MeanAggregator extends MultiRoundAggregator {
 
     this.log(AggregationStep.AGGREGATE);
 
-    const result = aggregation.avg(currentContributions.values());
+    const result = avg(currentContributions.values());
     return result;
   }
 
   override makePayloads(
     weights: WeightsContainer,
-  ): Map<client.NodeID, WeightsContainer> {
+  ): Map<NodeID, WeightsContainer> {
     // Communicate our local weights to every other node, be it a peer or a server
     return this.nodes.toMap().map(() => weights);
   }
diff --git a/discojs/src/aggregator/multiround.ts b/discojs/src/aggregator/multiround.ts
index 4018b55dd..d435cfaad 100644
--- a/discojs/src/aggregator/multiround.ts
+++ b/discojs/src/aggregator/multiround.ts
@@ -1,4 +1,4 @@
-import { Aggregator } from "./aggregator.js";
+import { Aggregator } from "#aggregator/aggregator";
 import createDebug from "debug";
 
 export type ThresholdType = "relative" | "absolute";
diff --git a/discojs/src/aggregator/percentile_clipping.spec.ts b/discojs/src/aggregator/percentile_clipping.spec.ts
index 755d5b382..39afa403b 100644
--- a/discojs/src/aggregator/percentile_clipping.spec.ts
+++ b/discojs/src/aggregator/percentile_clipping.spec.ts
@@ -1,8 +1,8 @@
 import { Set } from "immutable";
 import { describe, expect, it } from "vitest";
 
-import { WeightsContainer } from "../index.js";
-import { PercentileClippingAggregator } from "./percentile_clipping.js";
+import { WeightsContainer } from "#weights/index";
+import { PercentileClippingAggregator } from "#aggregator/percentile_clipping";
 
 async function WSIntoArrays(ws: WeightsContainer): Promise<number[][]> {
   return Promise.all(ws.weights.map(async (t) => Array.from(await t.data())));
diff --git a/discojs/src/aggregator/percentile_clipping.ts b/discojs/src/aggregator/percentile_clipping.ts
index 877a4bae3..5938c4cd8 100644
--- a/discojs/src/aggregator/percentile_clipping.ts
+++ b/discojs/src/aggregator/percentile_clipping.ts
@@ -1,9 +1,11 @@
-import { Map } from "immutable";
+import type { Map } from "immutable";
 import * as tf from "@tensorflow/tfjs";
-import { AggregationStep } from "./aggregator.js";
-import { MultiRoundAggregator, ThresholdType } from "./multiround.js";
-import { WeightsContainer, client } from "../index.js";
-import { aggregation } from "../index.js";
+import { AggregationStep } from "#aggregator/aggregator";
+import type { ThresholdType } from "#aggregator/multiround";
+import { MultiRoundAggregator } from "#aggregator/multiround";
+import type { NodeID } from "#client/index";
+import type { WeightsContainer } from "#weights/index";
+import { avg } from "#weights/index";
 
 /**
  * Percentile-based clipping aggregator.
@@ -51,7 +53,7 @@ export class PercentileClippingAggregator extends MultiRoundAggregator {
     this.tauPercentile = tauPercentile;
   }
 
-  override _add(nodeId: client.NodeID, contribution: WeightsContainer): void {
+  override _add(nodeId: NodeID, contribution: WeightsContainer): void {
     this.log(
       this.contributions.hasIn([0, nodeId])
         ? AggregationStep.UPDATE
@@ -74,9 +76,9 @@ export class PercentileClippingAggregator extends MultiRoundAggregator {
     if (this.prevAggregate) {
       centerReference = this.prevAggregate.map((t) => tf.clone(t));
     } else {
-      centerReference = aggregation
-        .avg(currentContributions.values())
-        .map((t) => tf.clone(t));
+      centerReference = avg(currentContributions.values()).map((t) =>
+        tf.clone(t),
+      );
     }
 
     // Step 2: Center the weights with respect to the reference
@@ -103,7 +105,7 @@ export class PercentileClippingAggregator extends MultiRoundAggregator {
     centeredWeights.forEach((w) => w.dispose());
 
     // Step 6: Average the clipped weights and add back the reference
-    const clippedAvg = aggregation.avg(clippedWeights);
+    const clippedAvg = avg(clippedWeights);
     const result = centerReference.add(clippedAvg);
 
     centerReference.dispose();
@@ -134,7 +136,7 @@ export class PercentileClippingAggregator extends MultiRoundAggregator {
 
   override makePayloads(
     weights: WeightsContainer,
-  ): Map<client.NodeID, WeightsContainer> {
+  ): Map<NodeID, WeightsContainer> {
     return this.nodes.toMap().map(() => weights);
   }
 }
diff --git a/discojs/src/aggregator/secure.spec.ts b/discojs/src/aggregator/secure.spec.ts
index 078658abe..d822bcb30 100644
--- a/discojs/src/aggregator/secure.spec.ts
+++ b/discojs/src/aggregator/secure.spec.ts
@@ -1,13 +1,10 @@
 import { List, Map, Range, Set } from "immutable";
 import { assert, describe, expect, it } from "vitest";
-import { communicate, setupNetwork, wsIntoArrays } from "../aggregator.spec.js";
-import {
-  WeightsContainer,
-  aggregation,
-  aggregator as aggregators,
-} from "../index.js";
-import { MeanAggregator } from "./mean.js";
-import { SecureAggregator } from "./secure.js";
+import { communicate, setupNetwork, wsIntoArrays } from "#root/aggregator.spec";
+import { sum, avg, WeightsContainer } from "#weights/index";
+
+import { MeanAggregator } from "#aggregator/mean";
+import { SecureAggregator } from "#aggregator/secure";
 
 describe("secret shares test", () => {
   const epsilon = 1e-4;
@@ -22,7 +19,7 @@ describe("secret shares test", () => {
   function buildShares(): List<List<WeightsContainer>> {
     const nodes = Set(secrets.keys()).map(String);
     return secrets.map((secret) => {
-      const aggregator = new aggregators.SecureAggregator();
+      const aggregator = new SecureAggregator();
       aggregator.setNodes(nodes);
       return aggregator.generateAllShares(secret);
     });
@@ -33,12 +30,12 @@ describe("secret shares test", () => {
   ): List<WeightsContainer> {
     return Range(0, secrets.size)
       .map((idx) => allShares.map((shares) => shares.get(idx)))
-      .map((shares) => aggregation.sum(shares as List<WeightsContainer>))
+      .map((shares) => sum(shares as List<WeightsContainer>))
       .toList();
   }
 
   it("recover secrets from shares", () => {
-    const recovered = buildShares().map((shares) => aggregation.sum(shares));
+    const recovered = buildShares().map((shares) => sum(shares));
     assert.isTrue(
       (
         recovered.zip(secrets) as List<[WeightsContainer, WeightsContainer]>
@@ -47,7 +44,7 @@ describe("secret shares test", () => {
   });
 
   it("derive aggregation result from partial sums", () => {
-    const actual = aggregation.avg(buildPartialSums(buildShares()));
+    const actual = avg(buildPartialSums(buildShares()));
     assert.isTrue(actual.equals(expected, epsilon));
   });
 });
diff --git a/discojs/src/aggregator/secure.ts b/discojs/src/aggregator/secure.ts
index c20221311..949fe8e8b 100644
--- a/discojs/src/aggregator/secure.ts
+++ b/discojs/src/aggregator/secure.ts
@@ -1,9 +1,11 @@
 import { Map, List, Range } from "immutable";
 import * as tf from "@tensorflow/tfjs";
 
-import { AggregationStep, Aggregator } from "./aggregator.js";
-import type { WeightsContainer, client } from "../index.js";
-import { aggregation } from "../index.js";
+import type { WeightsContainer } from "#weights/index";
+import type { NodeID } from "#client/types";
+import { sum, avg } from "#weights/index";
+
+import { AggregationStep, Aggregator } from "#aggregator/aggregator";
 
 /**
  * Aggregator implementing secure multi-party computation for decentralized learning.
@@ -27,7 +29,7 @@ export class SecureAggregator extends Aggregator {
         if (currentContributions === undefined)
           throw new Error("aggregating without any contribution");
 
-        return aggregation.sum(currentContributions.values());
+        return sum(currentContributions.values());
       }
       // Average the received partial sums
       case 1: {
@@ -35,7 +37,7 @@ export class SecureAggregator extends Aggregator {
         if (currentContributions === undefined)
           throw new Error("aggregating without any contribution");
 
-        return aggregation.avg(currentContributions.values());
+        return avg(currentContributions.values());
       }
       default:
         throw new Error("communication round is out of bounds");
@@ -43,7 +45,7 @@ export class SecureAggregator extends Aggregator {
   }
 
   _add(
-    nodeId: client.NodeID,
+    nodeId: NodeID,
     contribution: WeightsContainer,
     communicationRound: number,
   ): void {
@@ -77,7 +79,7 @@ export class SecureAggregator extends Aggregator {
 
   override makePayloads(
     weights: WeightsContainer,
-  ): Map<client.NodeID, WeightsContainer> {
+  ): Map<NodeID, WeightsContainer> {
     switch (this.communicationRound) {
       case 0: {
         const shares = this.generateAllShares(weights);
@@ -105,7 +107,7 @@ export class SecureAggregator extends Aggregator {
       .toList();
 
     // The last share completes the sum
-    return shares.push(secret.sub(aggregation.sum(shares)));
+    return shares.push(secret.sub(sum(shares)));
   }
 
   /**
diff --git a/discojs/src/aggregator/secure_history.spec.ts b/discojs/src/aggregator/secure_history.spec.ts
index 3426ad066..3b813096b 100644
--- a/discojs/src/aggregator/secure_history.spec.ts
+++ b/discojs/src/aggregator/secure_history.spec.ts
@@ -3,12 +3,12 @@ import { describe, expect, it, assert } from "vitest";
 
 import * as tf from "@tensorflow/tfjs";
 
-import { aggregation, WeightsContainer } from "../index.js";
+import { sum, avg, WeightsContainer } from "#weights/index";
 
-import { SecureHistoryAggregator } from "./secure_history.js";
-import { SecureAggregator } from "./secure.js";
+import { SecureHistoryAggregator } from "#aggregator/secure_history";
+import { SecureAggregator } from "#aggregator/secure";
 
-import { wsIntoArrays, communicate, setupNetwork } from "../aggregator.spec.js";
+import { wsIntoArrays, communicate, setupNetwork } from "#root/aggregator.spec";
 
 describe("Secure history aggregator", function () {
   const epsilon = 1e-4;
@@ -29,7 +29,7 @@ describe("Secure history aggregator", function () {
   }
 
   it("recovers secrets from shares", () => {
-    const recovered = buildShares().map((shares) => aggregation.sum(shares));
+    const recovered = buildShares().map((shares) => sum(shares));
     assert.isTrue(
       (
         recovered.zip(secrets) as List<[WeightsContainer, WeightsContainer]>
@@ -52,7 +52,7 @@ describe("Secure history aggregator", function () {
         const receivedShares = sharesRound0.map(
           (shares) => shares.get(receiverIdx)!,
         );
-        return aggregation.sum(receivedShares);
+        return sum(receivedShares);
       })
       .toList();
 
@@ -64,7 +64,7 @@ describe("Secure history aggregator", function () {
 
     const sumRound0 = await aggregationPromise;
 
-    const expectedSum = aggregation.sum(
+    const expectedSum = sum(
       sharesRound0.flatMap((x) => x), // flatten to List<WeightsContainer>
     );
     expect(sumRound0.equals(expectedSum, epsilon)).to.be.true;
@@ -79,7 +79,7 @@ describe("Secure history aggregator", function () {
     const sumRound1 = await aggregationPromise2;
 
     // First aggregation with momentum - no previous momentum, so just average
-    const avgPartialSum = aggregation.avg(partialSums);
+    const avgPartialSum = avg(partialSums);
     expect(sumRound1.equals(avgPartialSum, epsilon)).to.be.true;
 
     // Now we simulate a second round of aggregation with momentum smoothing
@@ -103,7 +103,7 @@ describe("Secure history aggregator", function () {
     });
     const sumRound2 = await aggregationPromise3;
 
-    const avgPartialSum2 = aggregation.avg(partialSums2);
+    const avgPartialSum2 = avg(partialSums2);
     const expectedSumRound2 = avgPartialSum.mapWith(
       avgPartialSum2,
       (prev, curr) => prev.mul(0.8).add(curr.mul(0.2)), // 0.8 = beta, 0.2 = (1 - beta)
diff --git a/discojs/src/aggregator/secure_history.ts b/discojs/src/aggregator/secure_history.ts
index 71c77b340..6b14fb9e7 100644
--- a/discojs/src/aggregator/secure_history.ts
+++ b/discojs/src/aggregator/secure_history.ts
@@ -1,6 +1,7 @@
-import type { WeightsContainer } from "../index.js";
-import { SecureAggregator } from "./secure.js";
-import { aggregation } from "../index.js";
+import type { WeightsContainer } from "#weights/index";
+import { avg } from "#weights/index";
+
+import { SecureAggregator } from "#aggregator/secure";
 
 /**
  * Aggregator that implements secure multi-party computation with history-based momentum smoothing.
@@ -43,15 +44,16 @@ export class SecureHistoryAggregator extends SecureAggregator {
     if (!currentContributions)
       throw new Error("aggregating without any contribution");
 
-    const avg = aggregation.avg(currentContributions.values());
+    const contribAvg = avg(currentContributions.values());
 
     if (this.prevAggregate === null) {
-      this.prevAggregate = avg;
-      return avg;
+      this.prevAggregate = contribAvg;
+      return contribAvg;
     }
 
-    const updatedMomentum = this.prevAggregate.mapWith(avg, (prevT, currT) =>
-      prevT.mul(this.beta).add(currT.mul(1 - this.beta)),
+    const updatedMomentum = this.prevAggregate.mapWith(
+      contribAvg,
+      (prevT, currT) => prevT.mul(this.beta).add(currT.mul(1 - this.beta)),
     );
 
     // Dispose old tensors to avoid memory leaks
diff --git a/discojs/src/client/client.ts b/discojs/src/client/client.ts
index 038d8ffad..1d23b559f 100644
--- a/discojs/src/client/client.ts
+++ b/discojs/src/client/client.ts
@@ -1,19 +1,18 @@
 import createDebug from "debug";
 
-import type {
-  DataType,
-  Model,
-  Network,
-  RoundStatus,
-  Task,
-  WeightsContainer,
-} from "../index.js";
-import { serialization } from "../index.js";
-import type { NodeID } from "./types.js";
-import type { EventConnection } from "./event_connection.js";
-import type { Aggregator } from "../aggregator/index.js";
-import { EventEmitter } from "../utils/event_emitter.js";
-import { type } from "./messages.js";
+import type { Model } from "#models/index";
+import type { DataType, Network } from "#types/index";
+import type { Task } from "#task/index";
+import type { WeightsContainer } from "#weights/index";
+import type { Aggregator } from "#aggregator/index";
+import type { RoundStatus } from "#training/types";
+import { EventEmitter } from "#utils/event_emitter";
+
+import { modelDecode } from "#serialization/index";
+
+import type { EventConnection } from "#client/event_connection";
+import type { NodeID } from "#client/types";
+import { MType } from "#client/mtype";
 
 const debug = createDebug("discojs:client");
 
@@ -120,7 +119,7 @@ export abstract class Client<N extends Network> extends EventEmitter<{
   protected setupServerCallbacks(setMessageInversionFlag: () => void) {
     // Setup an event callback if the server signals that we should
     // wait for more participants
-    this.server.on(type.WaitingForMoreParticipants, (event) => {
+    this.server.on(MType.WaitingForMoreParticipants, (event) => {
       if (this.promiseForMoreParticipants !== undefined)
         throw new Error(
           "Server sent multiple WaitingForMoreParticipants messages",
@@ -143,7 +142,7 @@ export abstract class Client<N extends Network> extends EventEmitter<{
     // and directly follows with an EnoughParticipants message when the 2nd participant joins
     // However, the EnoughParticipants can arrive before the NewNodeInfo (which can be much bigger)
     // so we check whether we received the EnoughParticipants before being assigned a node ID
-    this.server.once(type.EnoughParticipants, (event) => {
+    this.server.once(MType.EnoughParticipants, (event) => {
       if (this._ownId === undefined) {
         setMessageInversionFlag();
         this.nbOfParticipants = event.nbOfParticipants;
@@ -160,7 +159,7 @@ export abstract class Client<N extends Network> extends EventEmitter<{
   protected async createPromiseForMoreParticipants(): Promise<void> {
     return new Promise<void>((resolve) => {
       // "once" is important because we can't resolve the same promise multiple times
-      this.server.once(type.EnoughParticipants, (event) => {
+      this.server.once(MType.EnoughParticipants, (event) => {
         debug(
           `[${shortenId(this.ownId)}] received EnoughParticipants message from server`,
         );
@@ -201,7 +200,7 @@ export abstract class Client<N extends Network> extends EventEmitter<{
     if (!response.ok) throw new Error(`fetch: HTTP status ${response.status}`);
 
     const encoded = new Uint8Array(await response.arrayBuffer());
-    return await serialization.model.decode(encoded);
+    return await modelDecode(encoded);
   }
 
   /**
diff --git a/discojs/src/client/decentralized/decentralized_client.ts b/discojs/src/client/decentralized/decentralized_client.ts
index 474968102..85b7b0beb 100644
--- a/discojs/src/client/decentralized/decentralized_client.ts
+++ b/discojs/src/client/decentralized/decentralized_client.ts
@@ -1,20 +1,22 @@
 import createDebug from "debug";
 import { Map, Set } from "immutable";
 
-import type { DataType, Model, WeightsContainer } from "../../index.js";
-import { serialization } from "../../index.js";
-import { Client, shortenId } from "../client.js";
-import { type NodeID } from "../index.js";
-import { type, type ClientConnected } from "../messages.js";
-import { timeout } from "../utils.js";
+import type { WeightsContainer } from "#weights/index";
+import type { Model } from "#models/index";
+import type { DataType } from "#types/index";
+import { weightsEncode, weightsDecode } from "#serialization/index";
+import { Client, shortenId } from "#client/client";
+import type { NodeID } from "#client/types";
+import { MType, type ClientConnected } from "#client/mtype";
+import { timeout } from "#client/utils";
 import {
   WebSocketServer,
   waitMessage,
   type PeerConnection,
   waitMessageWithTimeout,
-} from "../event_connection.js";
-import { PeerPool } from "./peer_pool.js";
-import * as messages from "./messages.js";
+} from "#client/event_connection";
+import { PeerPool } from "#client/decentralized/peer_pool";
+import * as messages from "#client/decentralized/messages";
 
 const debug = createDebug("discojs:client:decentralized");
 
@@ -73,7 +75,7 @@ export class DecentralizedClient extends Client<"decentralized"> {
       messages.isMessageFromServer,
       messages.isMessageToServer,
     );
-    this.server.on(type.SignalForPeer, (event) => {
+    this.server.on(MType.SignalForPeer, (event) => {
       if (this.#pool === undefined)
         throw new Error("received signal but peer pool is undefined");
       // Create a WebRTC connection with the peer
@@ -85,13 +87,13 @@ export class DecentralizedClient extends Client<"decentralized"> {
     this.setupServerCallbacks(() => (receivedEnoughParticipants = true));
 
     const msg: ClientConnected = {
-      type: type.ClientConnected,
+      type: MType.ClientConnected,
     };
     this.server.send(msg);
 
     const { id, waitForMoreParticipants, nbOfParticipants } = await waitMessage(
       this.server,
-      type.NewDecentralizedNodeInfo,
+      MType.NewDecentralizedNodeInfo,
     );
 
     this.nbOfParticipants = nbOfParticipants;
@@ -145,7 +147,7 @@ export class DecentralizedClient extends Client<"decentralized"> {
   override async onRoundBeginCommunication(): Promise<void> {
     // Notify the server we want to join the next round so that the server
     // waits for us to be ready before sending the list of peers for the round
-    this.server.send({ type: type.JoinRound });
+    this.server.send({ type: MType.JoinRound });
     // Store the promise for the current round's aggregation result.
     // We will await for it to resolve at the end of the round when exchanging weight updates.
     this.aggregationResult = this.aggregator.getPromiseForAggregation();
@@ -190,7 +192,7 @@ export class DecentralizedClient extends Client<"decentralized"> {
     // Reset peers list at each round of training to make sure client works with an updated peers
     // list, maintained by the server. Adds any received weights to the aggregator.
     // Tell the server we are ready for the next round
-    const readyMessage: messages.PeerIsReady = { type: type.PeerIsReady };
+    const readyMessage: messages.PeerIsReady = { type: MType.PeerIsReady };
     this.server.send(readyMessage);
 
     // Wait for the server to answer with the list of peers for the round
@@ -200,7 +202,7 @@ export class DecentralizedClient extends Client<"decentralized"> {
       );
       const receivedMessage = await waitMessage(
         this.server,
-        type.PeersForRound,
+        MType.PeersForRound,
       );
 
       const peers = Set(receivedMessage.peers);
@@ -250,11 +252,11 @@ export class DecentralizedClient extends Client<"decentralized"> {
         try {
           const message = await waitMessageWithTimeout(
             connection,
-            type.Payload,
+            MType.Payload,
             60_000,
             "Timeout waiting for a contribution from peer " + peerId,
           );
-          const decoded = serialization.weights.decode(message.payload);
+          const decoded = weightsDecode(message.payload);
 
           if (
             !this.aggregator.isValidContribution(
@@ -335,9 +337,9 @@ export class DecentralizedClient extends Client<"decentralized"> {
         // Send our payload to each peer
         const peer = connections.get(id);
         if (peer !== undefined) {
-          const encoded = await serialization.weights.encode(payload);
+          const encoded = await weightsEncode(payload);
           const msg: messages.PeerMessage = {
-            type: type.Payload,
+            type: MType.Payload,
             peer: id,
             aggregationRound: this.aggregator.round,
             communicationRound,
diff --git a/discojs/src/client/decentralized/index.ts b/discojs/src/client/decentralized/index.ts
index 6e2c43bfa..5d2d88cb2 100644
--- a/discojs/src/client/decentralized/index.ts
+++ b/discojs/src/client/decentralized/index.ts
@@ -1,2 +1,3 @@
 export { DecentralizedClient } from "./decentralized_client.js";
+// eslint-disable-next-line no-restricted-syntax -- namespace re-export acceptable here
 export * as messages from "./messages.js";
diff --git a/discojs/src/client/decentralized/messages.ts b/discojs/src/client/decentralized/messages.ts
index 30991c3eb..9cfba5055 100644
--- a/discojs/src/client/decentralized/messages.ts
+++ b/discojs/src/client/decentralized/messages.ts
@@ -1,17 +1,17 @@
-import { serialization } from "../../index.js";
+import * as serialization from "#serialization/index";
 
-import { type SignalData } from "./peer.js";
-import { isNodeID, type NodeID } from "../types.js";
-import { type, hasMessageType } from "../messages.js";
+import { type SignalData } from "#client/decentralized/peer";
+import { isNodeID, type NodeID } from "#client/types";
+import { MType, hasMessageType } from "#client/mtype";
 import type {
   ClientConnected,
   WaitingForMoreParticipants,
   EnoughParticipants,
-} from "../messages.js";
+} from "#client/mtype";
 
 /// Phase 0 communication (between server and peers)
 export interface NewDecentralizedNodeInfo {
-  type: type.NewDecentralizedNodeInfo;
+  type: MType.NewDecentralizedNodeInfo;
   id: NodeID;
   waitForMoreParticipants: boolean;
   nbOfParticipants: number;
@@ -19,24 +19,24 @@ export interface NewDecentralizedNodeInfo {
 
 // WebRTC signal to forward to other node
 export interface SignalForPeer {
-  type: type.SignalForPeer;
+  type: MType.SignalForPeer;
   peer: NodeID;
   signal: SignalData;
 }
 
 // peer wants to join the next round
 export interface JoinRound {
-  type: type.JoinRound;
+  type: MType.JoinRound;
 }
 
 // peer who sent is ready
 export interface PeerIsReady {
-  type: type.PeerIsReady;
+  type: MType.PeerIsReady;
 }
 
 // server sends to each peer the list of peers to connect to
 export interface PeersForRound {
-  type: type.PeersForRound;
+  type: MType.PeersForRound;
   peers: NodeID[];
   aggregationRound: number;
 }
@@ -44,7 +44,7 @@ export interface PeersForRound {
 /// Phase 1 communication (between peers)
 
 export interface Payload {
-  type: type.Payload;
+  type: MType.Payload;
   peer: NodeID;
   aggregationRound: number;
   communicationRound: number;
@@ -72,19 +72,19 @@ export function isMessageFromServer(o: unknown): o is MessageFromServer {
   if (!hasMessageType(o)) return false;
 
   switch (o.type) {
-    case type.NewDecentralizedNodeInfo:
+    case MType.NewDecentralizedNodeInfo:
       return (
         "id" in o &&
         isNodeID(o.id) &&
         "waitForMoreParticipants" in o &&
         typeof o.waitForMoreParticipants === "boolean"
       );
-    case type.SignalForPeer:
+    case MType.SignalForPeer:
       return "peer" in o && isNodeID(o.peer) && "signal" in o; // TODO check signal content?
-    case type.PeersForRound:
+    case MType.PeersForRound:
       return "peers" in o && Array.isArray(o.peers) && o.peers.every(isNodeID);
-    case type.WaitingForMoreParticipants:
-    case type.EnoughParticipants:
+    case MType.WaitingForMoreParticipants:
+    case MType.EnoughParticipants:
       return true;
   }
 
@@ -95,12 +95,12 @@ export function isMessageToServer(o: unknown): o is MessageToServer {
   if (!hasMessageType(o)) return false;
 
   switch (o.type) {
-    case type.ClientConnected:
+    case MType.ClientConnected:
       return true;
-    case type.SignalForPeer:
+    case MType.SignalForPeer:
       return "peer" in o && isNodeID(o.peer) && "signal" in o; // TODO check signal content?
-    case type.JoinRound:
-    case type.PeerIsReady:
+    case MType.JoinRound:
+    case MType.PeerIsReady:
       return true;
   }
 
@@ -111,7 +111,7 @@ export function isPeerMessage(o: unknown): o is PeerMessage {
   if (!hasMessageType(o)) return false;
 
   switch (o.type) {
-    case type.Payload:
+    case MType.Payload:
       return (
         "peer" in o &&
         isNodeID(o.peer) &&
diff --git a/discojs/src/client/decentralized/peer.spec.ts b/discojs/src/client/decentralized/peer.spec.ts
index b089b43fb..89760fa5c 100644
--- a/discojs/src/client/decentralized/peer.spec.ts
+++ b/discojs/src/client/decentralized/peer.spec.ts
@@ -1,6 +1,6 @@
 import { List, Range, Set } from "immutable";
 import { assert, afterEach, beforeEach, describe, it } from "vitest";
-import { Peer } from "./peer.js";
+import { Peer } from "#client/decentralized/peer";
 
 describe("peer", () => {
   let peer1: Peer;
diff --git a/discojs/src/client/decentralized/peer.ts b/discojs/src/client/decentralized/peer.ts
index df8d41461..0ec23d77e 100644
--- a/discojs/src/client/decentralized/peer.ts
+++ b/discojs/src/client/decentralized/peer.ts
@@ -2,7 +2,7 @@ import { List, Map, Range, Seq } from "immutable";
 import wrtc from "@epfml/isomorphic-wrtc";
 import SimplePeer from "simple-peer";
 
-import type { NodeID } from "../types.js";
+import type { NodeID } from "#client/types";
 
 type MessageID = number;
 type ChunkID = number;
diff --git a/discojs/src/client/decentralized/peer_pool.spec.ts b/discojs/src/client/decentralized/peer_pool.spec.ts
index 435f7ef0e..7c519e557 100644
--- a/discojs/src/client/decentralized/peer_pool.spec.ts
+++ b/discojs/src/client/decentralized/peer_pool.spec.ts
@@ -1,12 +1,12 @@
 import { Map, Range } from "immutable";
 import { assert, afterEach, beforeEach, describe, it } from "vitest";
 
-import type { EventConnection, PeerConnection } from "../event_connection.js";
-import { type } from "../messages.js";
-import type { NodeID } from "../types.js";
+import type { EventConnection, PeerConnection } from "#client/event_connection";
+import { MType } from "#client/mtype";
+import type { NodeID } from "#client/types";
 
-import type { messages } from "./index.js";
-import { PeerPool } from "./peer_pool.js";
+import type * as messages from "#client/decentralized/messages";
+import { PeerPool } from "#client/decentralized/peer_pool";
 
 describe("peer pool", { timeout: 10_000 }, () => {
   let pools: Map<NodeID, PeerPool>;
@@ -47,7 +47,7 @@ describe("peer pool", { timeout: 10_000 }, () => {
 
   function mockWeights(id: NodeID): messages.Payload {
     return {
-      type: type.Payload,
+      type: MType.Payload,
       peer: id,
       payload: Uint8Array.of(1, 2, 3),
       aggregationRound: 0,
@@ -112,7 +112,7 @@ describe("peer pool", { timeout: 10_000 }, () => {
                   .map(
                     async (peer) =>
                       await new Promise<messages.Payload>((resolve) => {
-                        peer.on(type.Payload, (data) => {
+                        peer.on(MType.Payload, (data) => {
                           resolve(data);
                         });
                       }),
diff --git a/discojs/src/client/decentralized/peer_pool.ts b/discojs/src/client/decentralized/peer_pool.ts
index dd004b3a5..838916835 100644
--- a/discojs/src/client/decentralized/peer_pool.ts
+++ b/discojs/src/client/decentralized/peer_pool.ts
@@ -1,9 +1,9 @@
 import createDebug from "debug";
 import { Map, type Set } from "immutable";
 
-import { Peer, type SignalData } from "./peer.js";
-import type { NodeID } from "../types.js";
-import { PeerConnection, type EventConnection } from "../event_connection.js";
+import { Peer, type SignalData } from "#client/decentralized/peer";
+import type { NodeID } from "#client/types";
+import { PeerConnection, type EventConnection } from "#client/event_connection";
 
 const debug = createDebug("discojs:client:decentralized:pool");
 
diff --git a/discojs/src/client/event_connection.ts b/discojs/src/client/event_connection.ts
index ca253ba64..700b6994a 100644
--- a/discojs/src/client/event_connection.ts
+++ b/discojs/src/client/event_connection.ts
@@ -1,22 +1,23 @@
 import createDebug from "debug";
 import WebSocket from "isomorphic-ws";
 import * as msgpack from "@msgpack/msgpack";
-import type { Peer, SignalData } from "./decentralized/peer.js";
-import type { NodeID } from "./types.js";
-import * as decentralizedMessages from "./decentralized/messages.js";
-import { type, type NarrowMessage, type Message } from "./messages.js";
-import { timeout } from "./utils.js";
+import type { Peer, SignalData } from "#client/decentralized/peer";
+import type { NodeID } from "#client/types";
+import * as decentralizedMessages from "#client/decentralized/messages";
+import { MType } from "#client/mtype";
+import { type NarrowMessage, type Message } from "#client/messages";
+import { timeout } from "#client/utils";
 
-import { EventEmitter } from "../utils/event_emitter.js";
+import { EventEmitter } from "#utils/event_emitter";
 
 const debug = createDebug("discojs:client:connections");
 
 export interface EventConnection {
-  on: <K extends type>(
+  on: <K extends MType>(
     type: K,
     handler: (event: NarrowMessage<K>) => void,
   ) => void;
-  once: <K extends type>(
+  once: <K extends MType>(
     type: K,
     handler: (event: NarrowMessage<K>) => void,
   ) => void;
@@ -24,7 +25,7 @@ export interface EventConnection {
   disconnect: () => Promise<void>;
 }
 
-export async function waitMessage<T extends type>(
+export async function waitMessage<T extends MType>(
   connection: EventConnection,
   type: T,
 ): Promise<NarrowMessage<T>> {
@@ -36,7 +37,7 @@ export async function waitMessage<T extends type>(
   });
 }
 
-export async function waitMessageWithTimeout<T extends type>(
+export async function waitMessageWithTimeout<T extends MType>(
   connection: EventConnection,
   type: T,
   timeoutMs?: number,
@@ -49,7 +50,7 @@ export async function waitMessageWithTimeout<T extends type>(
 }
 
 export class PeerConnection
-  extends EventEmitter<{ [K in type]: NarrowMessage<K> }>
+  extends EventEmitter<{ [K in MType]: NarrowMessage<K> }>
   implements EventConnection
 {
   constructor(
@@ -63,7 +64,7 @@ export class PeerConnection
   async connect(): Promise<void> {
     this.peer.on("signal", (signal) => {
       const msg: decentralizedMessages.SignalForPeer = {
-        type: type.SignalForPeer,
+        type: MType.SignalForPeer,
         peer: this.peer.id,
         signal,
       };
@@ -108,7 +109,7 @@ export class PeerConnection
 }
 
 export class WebSocketServer
-  extends EventEmitter<{ [K in type]: NarrowMessage<K> }>
+  extends EventEmitter<{ [K in MType]: NarrowMessage<K> }>
   implements EventConnection
 {
   private constructor(
diff --git a/discojs/src/client/federated/federated_client.ts b/discojs/src/client/federated/federated_client.ts
index 9e21bd46a..d8812a622 100644
--- a/discojs/src/client/federated/federated_client.ts
+++ b/discojs/src/client/federated/federated_client.ts
@@ -1,11 +1,13 @@
 import createDebug from "debug";
 
-import { serialization } from "../../index.js";
-import type { DataType, Model, WeightsContainer } from "../../index.js";
-import { Client, shortenId } from "../client.js";
-import { type, type ClientConnected } from "../messages.js";
-import { waitMessage, WebSocketServer } from "../event_connection.js";
-import * as messages from "./messages.js";
+import type { Model } from "#models/index";
+import type { DataType } from "#types/index";
+import type { WeightsContainer } from "#weights/index";
+import { weightsEncode, weightsDecode } from "#serialization/index";
+import { Client, shortenId } from "#client/client";
+import { MType, type ClientConnected } from "#client/mtype";
+import { waitMessage, WebSocketServer } from "#client/event_connection";
+import * as messages from "#client/federated/messages";
 
 const debug = createDebug("discojs:client:federated");
 
@@ -55,12 +57,12 @@ export class FederatedClient extends Client<"federated"> {
     this.aggregator.registerNode(SERVER_NODE_ID);
 
     const msg: ClientConnected = {
-      type: type.ClientConnected,
+      type: MType.ClientConnected,
     };
     this.server.send(msg);
 
     const { id, waitForMoreParticipants, payload, round, nbOfParticipants } =
-      await waitMessage(this.server, type.NewFederatedNodeInfo);
+      await waitMessage(this.server, MType.NewFederatedNodeInfo);
 
     // This should come right after receiving the message to make sure
     // we don't miss a subsequent message from the server
@@ -85,7 +87,7 @@ export class FederatedClient extends Client<"federated"> {
       `[${shortenId(this.ownId)}] upon connecting, wait for participant flag %o`,
       this.waitingForMoreParticipants,
     );
-    model.weights = serialization.weights.decode(payload);
+    model.weights = weightsDecode(payload);
     return model;
   }
 
@@ -139,8 +141,8 @@ export class FederatedClient extends Client<"federated"> {
     if (payloadToServer === undefined)
       throw new Error("aggregator didn't make a payload for the server");
     const msg: messages.SendPayload = {
-      type: type.SendPayload,
-      payload: await serialization.weights.encode(payloadToServer),
+      type: MType.SendPayload,
+      payload: await weightsEncode(payloadToServer),
       round: this.aggregator.round,
     };
 
@@ -157,9 +159,9 @@ export class FederatedClient extends Client<"federated"> {
       payload: payloadFromServer,
       round: serverRound,
       nbOfParticipants,
-    } = await waitMessage(this.server, type.ReceiveServerPayload); // Wait indefinitely for the server update
+    } = await waitMessage(this.server, MType.ReceiveServerPayload); // Wait indefinitely for the server update
     this.nbOfParticipants = nbOfParticipants; // Save the current participants
-    const serverResult = serialization.weights.decode(payloadFromServer);
+    const serverResult = weightsDecode(payloadFromServer);
     this.aggregator.setRound(serverRound);
 
     return serverResult;
diff --git a/discojs/src/client/federated/index.ts b/discojs/src/client/federated/index.ts
index 056fd8f9f..ff6948724 100644
--- a/discojs/src/client/federated/index.ts
+++ b/discojs/src/client/federated/index.ts
@@ -1,2 +1,3 @@
 export { FederatedClient } from "./federated_client.js";
+// eslint-disable-next-line no-restricted-syntax -- namespace re-export acceptable here
 export * as messages from "./messages.js";
diff --git a/discojs/src/client/federated/messages.ts b/discojs/src/client/federated/messages.ts
index 4d4ee0e2a..d45e7c542 100644
--- a/discojs/src/client/federated/messages.ts
+++ b/discojs/src/client/federated/messages.ts
@@ -1,13 +1,12 @@
-import type { serialization } from "../../index.js";
+import type * as serialization from "#serialization/index";
+import type { NodeID } from "#client/types";
 
-import { type NodeID } from "..//types.js";
-
-import { type, hasMessageType } from "../messages.js";
+import { MType, hasMessageType } from "#client/mtype";
 import type {
   ClientConnected,
   WaitingForMoreParticipants,
   EnoughParticipants,
-} from "../messages.js";
+} from "#client/mtype";
 
 // See ../messages.ts for doc
 export type MessageFederated =
@@ -19,7 +18,7 @@ export type MessageFederated =
   | EnoughParticipants;
 
 export interface NewFederatedNodeInfo {
-  type: type.NewFederatedNodeInfo;
+  type: MType.NewFederatedNodeInfo;
   id: NodeID;
   waitForMoreParticipants: boolean;
   payload: serialization.Encoded;
@@ -28,12 +27,12 @@ export interface NewFederatedNodeInfo {
 }
 
 export interface SendPayload {
-  type: type.SendPayload;
+  type: MType.SendPayload;
   payload: serialization.Encoded;
   round: number;
 }
 export interface ReceiveServerPayload {
-  type: type.ReceiveServerPayload;
+  type: MType.ReceiveServerPayload;
   payload: serialization.Encoded;
   round: number;
   nbOfParticipants: number; // number of peers contributing to a federated training
@@ -45,12 +44,12 @@ export function isMessageFederated(raw: unknown): raw is MessageFederated {
   }
 
   switch (raw.type) {
-    case type.ClientConnected:
-    case type.NewFederatedNodeInfo:
-    case type.SendPayload:
-    case type.ReceiveServerPayload:
-    case type.WaitingForMoreParticipants:
-    case type.EnoughParticipants:
+    case MType.ClientConnected:
+    case MType.NewFederatedNodeInfo:
+    case MType.SendPayload:
+    case MType.ReceiveServerPayload:
+    case MType.WaitingForMoreParticipants:
+    case MType.EnoughParticipants:
       return true;
   }
 
diff --git a/discojs/src/client/get_client.ts b/discojs/src/client/get_client.ts
new file mode 100644
index 000000000..9678c8dbe
--- /dev/null
+++ b/discojs/src/client/get_client.ts
@@ -0,0 +1,41 @@
+import type { DataType, Network } from "#types/index";
+import type { Task } from "#task/index";
+import type * as aggregator from "#aggregator/index";
+
+// import * as clients from "#client/index";
+import { LocalClient } from "#client/local_client";
+import type { Client } from "#client/client";
+import { DecentralizedClient } from "#client/decentralized/decentralized_client";
+import { FederatedClient } from "#client/federated/federated_client";
+
+export function getClient<D extends DataType, N extends Network>(
+  scheme: N | "local",
+  serverURL: URL,
+  task: Task<D, N>,
+  aggregator: aggregator.Aggregator,
+): Client<N> {
+  switch (scheme) {
+    case "decentralized": {
+      const t = task as Task<D, "decentralized">;
+      t.trainingInformation.scheme = scheme;
+
+      return new DecentralizedClient(serverURL, t, aggregator);
+    }
+    case "federated": {
+      const t = task as Task<D, "federated">;
+      t.trainingInformation.scheme = scheme;
+
+      return new FederatedClient(serverURL, t, aggregator);
+    }
+    case "local": {
+      const t = task as Task<D, "local">;
+      t.trainingInformation.scheme = scheme;
+
+      return new LocalClient(serverURL, t, aggregator);
+    }
+    default: {
+      const _: never = scheme;
+      throw new Error("should never happen");
+    }
+  }
+}
diff --git a/discojs/src/client/index.ts b/discojs/src/client/index.ts
index f084d5a2f..bad8f14f0 100644
--- a/discojs/src/client/index.ts
+++ b/discojs/src/client/index.ts
@@ -1,11 +1,18 @@
 export { Client } from "./client.js";
 
-export * from "./types.js";
+export type { NodeID } from "./types.js";
 
-export * as aggregator from "../aggregator/index.js";
-export * as decentralized from "./decentralized/index.js";
-export * as federated from "./federated/index.js";
-export * as messages from "./messages.js";
-export { getClient, timeout } from "./utils.js";
+export {
+  messages as decentralizedMessages,
+  DecentralizedClient,
+} from "./decentralized/index.js";
+export {
+  messages as federatedMessages,
+  FederatedClient,
+} from "./federated/index.js";
+//
+// eslint-disable-next-line no-restricted-syntax -- namespace re-export acceptable here
+export * as mtype from "./mtype.js";
+export { getClient } from "./get_client.js";
 
 export { LocalClient } from "./local_client.js";
diff --git a/discojs/src/client/local_client.ts b/discojs/src/client/local_client.ts
index 1b40cc6bb..24e4eb59e 100644
--- a/discojs/src/client/local_client.ts
+++ b/discojs/src/client/local_client.ts
@@ -1,5 +1,5 @@
-import { WeightsContainer } from "../index.js";
-import { Client } from "./client.js";
+import type { WeightsContainer } from "#weights/index";
+import { Client } from "#client/client";
 
 /**
  * A LocalClient represents a Disco user training only on their local data without collaborating
diff --git a/discojs/src/client/messages.ts b/discojs/src/client/messages.ts
index afc2f6556..db7002b44 100644
--- a/discojs/src/client/messages.ts
+++ b/discojs/src/client/messages.ts
@@ -1,59 +1,5 @@
-import type * as decentralized from "./decentralized/messages.js";
-import type * as federated from "./federated/messages.js";
-
-export enum type {
-  // Sent from client to server as first point of contact to join a task.
-  // The server answers with an node id in a NewFederatedNodeInfo
-  // or NewDecentralizedNodeInfo message
-  ClientConnected,
-
-  /* Decentralized */
-  // When a user joins a task with a ClientConnected message, the server
-  // answers with  its peer id and also tells the client whether we are waiting
-  // for more participants before starting training
-  NewDecentralizedNodeInfo,
-  // Message sent by peers to the server to signal they want to
-  // join the next round
-  JoinRound,
-  // Message sent by nodes to server signaling they are ready to
-  // start the next round
-  PeerIsReady,
-  // Sent by the server to participating peers containing the list
-  // of peers for the round
-  PeersForRound,
-  // Message forwarded by the server from a client to another client
-  // to establish a peer-to-peer (WebRTC) connection
-  SignalForPeer,
-  // The weight update
-  Payload,
-
-  /* Federated */
-  // The server answers the ClientConnected message with the necessary information
-  // to start training: node id, latest model global weights, current round etc
-  NewFederatedNodeInfo,
-  // Message sent by server to notify clients that there are not enough
-  // participants to continue training
-  WaitingForMoreParticipants,
-  // Message sent by server to notify clients that there are now enough
-  // participants to start training collaboratively
-  EnoughParticipants,
-  SendPayload,
-  ReceiveServerPayload,
-}
-
-export interface ClientConnected {
-  type: type.ClientConnected;
-}
-
-export interface EnoughParticipants {
-  type: type.EnoughParticipants;
-  nbOfParticipants: number;
-}
-
-export interface WaitingForMoreParticipants {
-  type: type.WaitingForMoreParticipants;
-  nbOfParticipants: number;
-}
+import type * as decentralized from "#client/decentralized/messages";
+import type * as federated from "#client/federated/messages";
 
 export type Message =
   | decentralized.MessageFromServer
@@ -63,16 +9,3 @@ export type Message =
 
 // Retrieve a specific message interface from the type D. i.e. NarrowMessage<messages.type.PeerId> => messages.PeerId type
 export type NarrowMessage<D> = Extract<Message, { type: D }>;
-
-export function hasMessageType(
-  raw: unknown,
-): raw is { type: type } & Record<string, unknown> {
-  if (typeof raw !== "object" || raw === null) return false;
-
-  const o = raw as Record<string, unknown>;
-  if (!("type" in o && typeof o.type === "number" && o.type in type)) {
-    return false;
-  }
-
-  return true;
-}
diff --git a/discojs/src/client/mtype.ts b/discojs/src/client/mtype.ts
new file mode 100644
index 000000000..f271cbe08
--- /dev/null
+++ b/discojs/src/client/mtype.ts
@@ -0,0 +1,66 @@
+export enum MType {
+  // Sent from client to server as first point of contact to join a task.
+  // The server answers with an node id in a NewFederatedNodeInfo
+  // or NewDecentralizedNodeInfo message
+  ClientConnected,
+
+  /* Decentralized */
+  // When a user joins a task with a ClientConnected message, the server
+  // answers with  its peer id and also tells the client whether we are waiting
+  // for more participants before starting training
+  NewDecentralizedNodeInfo,
+  // Message sent by peers to the server to signal they want to
+  // join the next round
+  JoinRound,
+  // Message sent by nodes to server signaling they are ready to
+  // start the next round
+  PeerIsReady,
+  // Sent by the server to participating peers containing the list
+  // of peers for the round
+  PeersForRound,
+  // Message forwarded by the server from a client to another client
+  // to establish a peer-to-peer (WebRTC) connection
+  SignalForPeer,
+  // The weight update
+  Payload,
+
+  /* Federated */
+  // The server answers the ClientConnected message with the necessary information
+  // to start training: node id, latest model global weights, current round etc
+  NewFederatedNodeInfo,
+  // Message sent by server to notify clients that there are not enough
+  // participants to continue training
+  WaitingForMoreParticipants,
+  // Message sent by server to notify clients that there are now enough
+  // participants to start training collaboratively
+  EnoughParticipants,
+  SendPayload,
+  ReceiveServerPayload,
+}
+
+export function hasMessageType(
+  raw: unknown,
+): raw is { type: MType } & Record<string, unknown> {
+  if (typeof raw !== "object" || raw === null) return false;
+
+  const o = raw as Record<string, unknown>;
+  if (!("type" in o && typeof o.type === "number" && o.type in MType)) {
+    return false;
+  }
+
+  return true;
+}
+
+export interface ClientConnected {
+  type: MType.ClientConnected;
+}
+
+export interface EnoughParticipants {
+  type: MType.EnoughParticipants;
+  nbOfParticipants: number;
+}
+
+export interface WaitingForMoreParticipants {
+  type: MType.WaitingForMoreParticipants;
+  nbOfParticipants: number;
+}
diff --git a/discojs/src/client/utils.ts b/discojs/src/client/utils.ts
index cfc3e6922..20a5d8e5c 100644
--- a/discojs/src/client/utils.ts
+++ b/discojs/src/client/utils.ts
@@ -1,6 +1,3 @@
-import type { DataType, Network, Task } from "../index.js";
-import { client as clients, type aggregator } from "../index.js";
-
 // Time to wait for the others in milliseconds.
 const MAX_WAIT_PER_ROUND = 15_000;
 
@@ -14,39 +11,3 @@ export async function timeout(
     }, ms);
   });
 }
-
-export function getClient<D extends DataType, N extends Network>(
-  scheme: N | "local",
-  serverURL: URL,
-  task: Task<D, N>,
-  aggregator: aggregator.Aggregator,
-): clients.Client<N> {
-  switch (scheme) {
-    case "decentralized": {
-      const t = task as Task<D, "decentralized">;
-      t.trainingInformation.scheme = scheme;
-
-      return new clients.decentralized.DecentralizedClient(
-        serverURL,
-        t,
-        aggregator,
-      );
-    }
-    case "federated": {
-      const t = task as Task<D, "federated">;
-      t.trainingInformation.scheme = scheme;
-
-      return new clients.federated.FederatedClient(serverURL, t, aggregator);
-    }
-    case "local": {
-      const t = task as Task<D, "local">;
-      t.trainingInformation.scheme = scheme;
-
-      return new clients.LocalClient(serverURL, t, aggregator);
-    }
-    default: {
-      const _: never = scheme;
-      throw new Error("should never happen");
-    }
-  }
-}
diff --git a/discojs/src/dataset/dataset.spec.ts b/discojs/src/dataset/dataset.spec.ts
index 65c91e937..5ea6646ed 100644
--- a/discojs/src/dataset/dataset.spec.ts
+++ b/discojs/src/dataset/dataset.spec.ts
@@ -1,6 +1,6 @@
 import { List, Range } from "immutable";
 import { describe, expect, it } from "vitest";
-import { Dataset } from "./dataset.js";
+import { Dataset } from "#dataset/dataset";
 
 // Array.fromAsync not yet widely used (2024)
 async function arrayFromAsync<T>(iter: AsyncIterable<T>): Promise<T[]> {
diff --git a/discojs/src/dataset/dataset.ts b/discojs/src/dataset/dataset.ts
index be7085ef8..03b1f883b 100644
--- a/discojs/src/dataset/dataset.ts
+++ b/discojs/src/dataset/dataset.ts
@@ -1,7 +1,7 @@
 import createDebug from "debug";
 import { List, Range } from "immutable";
 
-import { Batched } from "./types.js";
+import type { Batched } from "#dataset/types";
 
 const debug = createDebug("discojs:dataset");
 
diff --git a/discojs/src/dataset/index.ts b/discojs/src/dataset/index.ts
index 3214f4050..4bb673edd 100644
--- a/discojs/src/dataset/index.ts
+++ b/discojs/src/dataset/index.ts
@@ -1,2 +1,3 @@
 export { Dataset } from "./dataset.js";
-export * from "./types.js";
+export { Image } from "./types.js";
+export type { Batched, Tabular, Text, TokenizedText } from "./types.js";
diff --git a/discojs/src/dataset/types.ts b/discojs/src/dataset/types.ts
index 915cedc71..c6091fed6 100644
--- a/discojs/src/dataset/types.ts
+++ b/discojs/src/dataset/types.ts
@@ -1,10 +1,9 @@
-import { List } from "immutable";
+import type { List } from "immutable";
 
-import { Image } from "./image.js";
+export { Image } from "#dataset/image";
 
 export type Batched<T> = List<T>;
 
-export { Image };
 export type Tabular = Partial<Record<string, string>>;
 export type Text = string;
 export type TokenizedText = List<number>;
diff --git a/discojs/src/index.ts b/discojs/src/index.ts
index 57a4711f0..fd87f7752 100644
--- a/discojs/src/index.ts
+++ b/discojs/src/index.ts
@@ -1,41 +1,78 @@
-export * as data from "./dataset/index.js";
-export * as serialization from "./serialization/index.js";
-export * as training from "./training/index.js";
-export * as privacy from "./privacy.js";
+export {
+  modelEncode,
+  modelDecode,
+  weightsEncode,
+  weightsDecode,
+  serializeTaskToJSON,
+  deserializeTaskFromJSON,
+  isEncoded,
+} from "./serialization/index.js";
+export type { Encoded } from "./serialization/index.js";
 
-export * as client from "./client/index.js";
-export * as aggregator from "./aggregator/index.js";
+export {
+  MeanAggregator,
+  SecureAggregator,
+  getAggregator,
+} from "./aggregator/index.js";
 
-export { WeightsContainer, aggregation } from "./weights/index.js";
-export { Logger, ConsoleLogger } from "./logging/index.js";
 export {
-  Disco,
-  RoundLogs,
-  RoundStatus,
-  SummaryLogs,
-} from "./training/index.js";
+  LocalClient,
+  getClient,
+  DecentralizedClient,
+  FederatedClient,
+  mtype,
+  federatedMessages,
+  decentralizedMessages,
+} from "./client/index.js";
+export type { Client, NodeID } from "./client/index.js";
+
+export { WeightsContainer, avg } from "./weights/index.js";
+
+export { Disco } from "./training/index.js";
+export type { RoundLogs, RoundStatus, SummaryLogs } from "./training/index.js";
+
 export { Validator } from "./validator.js";
 
-export {
-  Model,
+export type {
   ModelCard,
   ModelCardInfo,
   BatchLogs,
+  HellaSwagDataset,
+} from "./models/index.js";
+
+export {
+  Model,
   EpochLogs,
   Tokenizer,
-  ValidationMetrics,
   fetchModels,
+  GPT,
+  TFJS,
+  ONNXModel,
+  HELLASWAG_URL,
+  evaluate_hellaswag,
 } from "./models/index.js";
-export * as models from "./models/index.js";
-
-export * from "./task/index.js";
-export * as defaultTasks from "./default_tasks/index.js";
-export * as defaultModels from "./models/cards/index.js";
+export type { GPTConfig, HellaSwagExample } from "./models/index.js";
 
-export * as async_iterator from "./utils/async_iterator.js";
 export { EventEmitter } from "./utils/event_emitter.js";
 
-export * from "./dataset/index.js";
-export * from "./types/index.js";
+export { Dataset, Image } from "./dataset/index.js";
+export type { Text, Tabular } from "./dataset/index.js";
 
-export * as processing from "./processing/index.js";
+export { split, gather } from "./utils/async_iterator.js";
+
+export {
+  Task,
+  TrainingInformation,
+  pushTask,
+  fetchTasks,
+} from "./task/index.js";
+export type { TaskProvider } from "./task/index.js";
+
+export type { DataType, Network, DataFormat } from "./types/index.js";
+export { dataTypeValues } from "./types/index.js";
+
+export { extractColumn } from "./processing/index.js";
+
+// eslint-disable-next-line no-restricted-syntax -- namespace re-export acceptable here
+export * as defaultTasks from "./default_tasks/index.js";
+export { cards as defaultModels } from "./models/index.js";
diff --git a/discojs/src/logging/console_logger.ts b/discojs/src/logging/console_logger.ts
index de83dfcef..7c8e6e86b 100644
--- a/discojs/src/logging/console_logger.ts
+++ b/discojs/src/logging/console_logger.ts
@@ -1,5 +1,5 @@
 import chalk from "chalk";
-import { Logger } from "./logger.js";
+import type { Logger } from "#logging/logger";
 
 /**
  * Same properties as Toaster but on the console
diff --git a/discojs/src/logging/index.ts b/discojs/src/logging/index.ts
index 5b43b5751..397a45f31 100644
--- a/discojs/src/logging/index.ts
+++ b/discojs/src/logging/index.ts
@@ -1,2 +1,2 @@
-export { Logger } from "./logger.js";
+export type { Logger } from "./logger.js";
 export { ConsoleLogger } from "./console_logger.js";
diff --git a/discojs/src/models/cards/CIFAR10Classifier.ts b/discojs/src/models/cards/CIFAR10Classifier.ts
index c59b7c2eb..e198c7b89 100644
--- a/discojs/src/models/cards/CIFAR10Classifier.ts
+++ b/discojs/src/models/cards/CIFAR10Classifier.ts
@@ -1,4 +1,5 @@
-import { Model, ModelCard } from "#models/index";
+import type { Model } from "#models/model";
+import type { ModelCard } from "#models/model_card";
 import { getModel } from "#models/implementations/CIFAR10ClassifierModel";
 
 export const CIFAR10Classifier: ModelCard<"image"> = {
diff --git a/discojs/src/models/cards/LUSClassifier.ts b/discojs/src/models/cards/LUSClassifier.ts
index 0630adfbd..bbc0f0de0 100644
--- a/discojs/src/models/cards/LUSClassifier.ts
+++ b/discojs/src/models/cards/LUSClassifier.ts
@@ -1,5 +1,5 @@
-import { Model } from "#models/model";
-import { ModelCard } from "#models/model_card";
+import type { Model } from "#models/model";
+import type { ModelCard } from "#models/model_card";
 import { model } from "#models/implementations/LUSClassifierModel";
 
 export const LUSClassifier: ModelCard<"image"> = {
diff --git a/discojs/src/models/cards/MNISTClassifier.ts b/discojs/src/models/cards/MNISTClassifier.ts
index 0bb6c4ca8..96f911f45 100644
--- a/discojs/src/models/cards/MNISTClassifier.ts
+++ b/discojs/src/models/cards/MNISTClassifier.ts
@@ -1,5 +1,6 @@
-import { Model, ModelCard } from "../index.js";
-import { model } from "../implementations/MNISTClassifierModel.js";
+import type { Model } from "#models/model";
+import type { ModelCard } from "#models/model_card";
+import { model } from "#models/implementations/MNISTClassifierModel";
 
 export const MNISTClassifier: ModelCard<"image"> = {
   card: {
diff --git a/discojs/src/models/cards/dogClassifier.ts b/discojs/src/models/cards/dogClassifier.ts
index 6e7c762c8..a140e3865 100644
--- a/discojs/src/models/cards/dogClassifier.ts
+++ b/discojs/src/models/cards/dogClassifier.ts
@@ -1,5 +1,6 @@
-import { Model, ModelCard } from "../index.js";
-import { model } from "../implementations/dogClassifierModel.js";
+import type { Model } from "#models/model";
+import type { ModelCard } from "#models/model_card";
+import { model } from "#models/implementations/dogClassifierModel";
 
 export const DogClassifier: ModelCard<"image"> = {
   card: {
diff --git a/discojs/src/models/cards/titanicClassifier.ts b/discojs/src/models/cards/titanicClassifier.ts
index a4b33e209..436639491 100644
--- a/discojs/src/models/cards/titanicClassifier.ts
+++ b/discojs/src/models/cards/titanicClassifier.ts
@@ -1,5 +1,6 @@
-import { Model, ModelCard } from "../index.js";
-import { model } from "../implementations/titanicClassifierModel.js";
+import type { Model } from "#models/model";
+import type { ModelCard } from "#models/model_card";
+import { model } from "#models/implementations/titanicClassifierModel";
 
 export const TitanicClassifier: ModelCard<"tabular"> = {
   card: {
diff --git a/discojs/src/models/cards/wikitext.ts b/discojs/src/models/cards/wikitext.ts
index c452693dd..a698ed5b7 100644
--- a/discojs/src/models/cards/wikitext.ts
+++ b/discojs/src/models/cards/wikitext.ts
@@ -1,5 +1,6 @@
-import { Model, ModelCard } from "../index.js";
-import { GPT } from "../index.js";
+import type { Model } from "#models/model";
+import type { ModelCard } from "#models/model_card";
+import { GPT } from "#models/implementations/index";
 
 export const Wikitext: ModelCard<"text"> = {
   card: {
diff --git a/discojs/src/models/implementations/CIFAR10ClassifierModel.ts b/discojs/src/models/implementations/CIFAR10ClassifierModel.ts
index dfb5dd1a8..0435ce453 100644
--- a/discojs/src/models/implementations/CIFAR10ClassifierModel.ts
+++ b/discojs/src/models/implementations/CIFAR10ClassifierModel.ts
@@ -1,8 +1,8 @@
 import * as tf from "@tensorflow/tfjs";
 
-import { TFJS } from "../index.js";
+import { TFJS } from "#models/tfjs";
 
-import baseModel from "./mobileNet_v1_025_224.js";
+import baseModel from "#models/implementations/mobileNet_v1_025_224";
 
 export async function getModel() {
   const mobilenet = await tf.loadLayersModel({
diff --git a/discojs/src/models/implementations/LUSClassifierModel.ts b/discojs/src/models/implementations/LUSClassifierModel.ts
index 96174b58b..2d34cd76b 100644
--- a/discojs/src/models/implementations/LUSClassifierModel.ts
+++ b/discojs/src/models/implementations/LUSClassifierModel.ts
@@ -1,6 +1,6 @@
 import * as tf from "@tensorflow/tfjs";
 
-import { Model } from "#models/model";
+import type { Model } from "#models/model";
 import { TFJS } from "#models/tfjs";
 
 // Model architecture from tensorflow.js docs:
diff --git a/discojs/src/models/implementations/MNISTClassifierModel.ts b/discojs/src/models/implementations/MNISTClassifierModel.ts
index 582874f8f..27104f60e 100644
--- a/discojs/src/models/implementations/MNISTClassifierModel.ts
+++ b/discojs/src/models/implementations/MNISTClassifierModel.ts
@@ -1,6 +1,6 @@
 import * as tf from "@tensorflow/tfjs";
 
-import { TFJS } from "../index.js";
+import { TFJS } from "#models/tfjs";
 
 export function model() {
   // Architecture from the PyTorch MNIST example (I made it slightly smaller, 650kB instead of 5MB)
diff --git a/discojs/src/models/implementations/dogClassifierModel.ts b/discojs/src/models/implementations/dogClassifierModel.ts
index ff66091d9..ce790b4de 100644
--- a/discojs/src/models/implementations/dogClassifierModel.ts
+++ b/discojs/src/models/implementations/dogClassifierModel.ts
@@ -1,6 +1,6 @@
 import * as tf from "@tensorflow/tfjs";
 
-import { TFJS } from "../index.js";
+import { TFJS } from "#models/tfjs";
 
 export function model() {
   const seed = 42; // set a seed to ensure reproducibility during GDHF demo
diff --git a/discojs/src/models/implementations/gpt/gpt.spec.ts b/discojs/src/models/implementations/gpt/gpt.spec.ts
index c7e775545..467110e03 100644
--- a/discojs/src/models/implementations/gpt/gpt.spec.ts
+++ b/discojs/src/models/implementations/gpt/gpt.spec.ts
@@ -1,10 +1,11 @@
 import { List } from "immutable";
 import { describe, expect, it } from "vitest";
 
-import type { DataFormat } from "../../../index.js";
-import { Dataset, Tokenizer } from "../../../index.js";
+import type { DataFormat } from "#types/index";
+import { Dataset } from "#dataset/index";
 
-import { GPT } from "./index.js";
+import { Tokenizer } from "#models/tokenizer";
+import { GPT } from "#models/implementations/gpt/index";
 
 describe("gpt-tfjs", () => {
   it("can overfit one sentence", { timeout: 100_000 }, async () => {
diff --git a/discojs/src/models/implementations/gpt/gpt.ts b/discojs/src/models/implementations/gpt/gpt.ts
new file mode 100644
index 000000000..ed4cdb2ca
--- /dev/null
+++ b/discojs/src/models/implementations/gpt/gpt.ts
@@ -0,0 +1,274 @@
+/**
+ * Source: https://github.com/zemlyansky/gpt-tfjs and https://github.com/karpathy/build-nanogpt
+ * With modifications from @peacefulotter, @lukemovement and the Disco team
+ **/
+import createDebug from "debug";
+import { List, Range } from "immutable";
+import * as tf from "@tensorflow/tfjs";
+
+import { WeightsContainer } from "#weights/index";
+import type { Dataset, Batched } from "#dataset/index";
+import type { DataFormat } from "#types/index";
+
+import type { BatchLogs } from "#models/logs";
+import { EpochLogs } from "#models/logs";
+import { Model } from "#models/model";
+import { GPTModel } from "#models/implementations/gpt/model";
+import evaluate from "#models/implementations/gpt/evaluate";
+import {
+  DefaultGPTConfig,
+  DefaultGenerationConfig,
+} from "#models/implementations/gpt/config";
+import type {
+  GPTConfig,
+  GenerationConfig,
+} from "#models/implementations/gpt/config";
+
+const debug = createDebug("discojs:models:gpt");
+
+export type GPTSerialization = {
+  weights: WeightsContainer;
+  config?: GPTConfig;
+};
+
+export class GPT extends Model<"text"> {
+  readonly datatype = "text" as const;
+  private readonly model: GPTModel;
+
+  readonly #contextLength: number;
+  readonly #maxBatchCount: number;
+  readonly #vocabSize: number;
+
+  constructor(
+    partialConfig?: Partial<GPTConfig>,
+    layersModel?: tf.LayersModel,
+  ) {
+    super();
+
+    const model = new GPTModel(partialConfig, layersModel);
+    model.compile();
+    this.model = model;
+
+    this.#contextLength =
+      partialConfig?.contextLength ?? DefaultGPTConfig.contextLength;
+    this.#maxBatchCount = partialConfig?.maxIter ?? DefaultGPTConfig.maxIter;
+    this.#vocabSize = partialConfig?.vocabSize ?? DefaultGPTConfig.vocabSize;
+  }
+
+  /**
+   * The GPT train methods wraps the model.fitDataset call in a for loop to act as a generator (of logs)
+   * This allows for getting logs and stopping training without callbacks.
+   *
+   * @param trainingData training dataset
+   * @param validationData validation dataset
+   * @param epochs the number of passes of the training dataset
+   * @param tracker
+   */
+  override async *train(
+    trainingDataset: Dataset<Batched<DataFormat.ModelEncoded["text"]>>,
+    validationDataset?: Dataset<Batched<DataFormat.ModelEncoded["text"]>>,
+  ): AsyncGenerator<BatchLogs, EpochLogs> {
+    let batchesLogs = List<BatchLogs>();
+    let epochTime = performance.now();
+
+    for await (const [batch, _] of trainingDataset.zip(
+      Range(0, this.#maxBatchCount),
+    )) {
+      const batchLogs = await this.#runBatch(batch);
+
+      yield batchLogs;
+      batchesLogs = batchesLogs.push(batchLogs);
+    }
+
+    const validation =
+      validationDataset && (await this.evaluate(validationDataset));
+    epochTime = performance.now() - epochTime;
+
+    return new EpochLogs(batchesLogs, epochTime, validation);
+  }
+
+  async #runBatch(
+    batch: Batched<DataFormat.ModelEncoded["text"]>,
+  ): Promise<BatchLogs> {
+    const tfBatch = this.#batchToTF(batch);
+
+    let logs: tf.Logs | undefined;
+    await this.model.fitDataset(tf.data.array([tfBatch]), {
+      epochs: 1,
+      verbose: 0, // don't pollute
+      callbacks: {
+        onEpochEnd: (_, cur) => {
+          logs = cur;
+        },
+      },
+    });
+    tf.dispose(tfBatch);
+    if (logs === undefined) throw new Error("batch didn't gave any logs");
+
+    const { loss, acc: accuracy } = logs;
+    if (loss === undefined || isNaN(loss))
+      throw new Error("training loss is undefined or NaN");
+
+    return {
+      accuracy,
+      loss,
+      memoryUsage: tf.memory().numBytes / 1024 / 1024 / 1024,
+    };
+  }
+
+  override async evaluate(
+    dataset: Dataset<Batched<DataFormat.ModelEncoded["text"]>>,
+  ): Promise<Record<"accuracy" | "loss", number>> {
+    const evaluation = await evaluate(
+      this.model,
+      tf.data.generator(
+        async function* (this: GPT) {
+          yield* dataset.map((batch) => this.#batchToTF(batch));
+        }.bind(this),
+      ),
+      this.config.maxEvalBatches,
+    );
+
+    return {
+      accuracy: evaluation.val_acc,
+      loss: evaluation.val_loss,
+    };
+  }
+
+  #batchToTF(batch: Batched<DataFormat.ModelEncoded["text"]>): {
+    xs: tf.Tensor2D;
+    ys: tf.Tensor3D;
+  } {
+    return tf.tidy(() => ({
+      xs: tf.stack(
+        batch.map(([line]) => tf.tensor1d(line.toArray(), "int32")).toArray(),
+      ) as tf.Tensor2D, // cast as stack doesn't type
+      ys: tf.stack(
+        batch
+          .map(([line, next]) =>
+            tf.oneHot(line.shift().push(next).toArray(), this.#vocabSize),
+          )
+          .toArray(),
+      ) as tf.Tensor3D, // cast as oneHot/stack doesn't type
+    }));
+  }
+
+  override async predict(
+    batch: Batched<DataFormat.ModelEncoded["text"][0]>,
+    options?: Partial<GenerationConfig>,
+  ): Promise<Batched<DataFormat.ModelEncoded["text"][1]>> {
+    // overwrite default with user config
+    const config = Object.assign({}, DefaultGenerationConfig, options);
+
+    return List(
+      await Promise.all(
+        batch.map((tokens) => this.#predictSingle(tokens, config)),
+      ),
+    );
+  }
+
+  /**
+   * Generate the next token after the input sequence.
+   * In other words, takes an input tensor of shape (prompt length T) and returns a tensor of shape (T+1)
+   *
+   * @param token input tokens of shape (T,). T is truncated to the model's context length
+   * @param config generation config: temperature, doSample, topk
+   * @returns the next token predicted by the model
+   */
+  async #predictSingle(
+    tokens: DataFormat.ModelEncoded["text"][0],
+    config: GenerationConfig,
+  ): Promise<DataFormat.ModelEncoded["text"][1]> {
+    // slice input tokens if longer than context length
+    tokens = tokens.slice(-this.#contextLength);
+
+    const input = tf.tidy(() =>
+      tf.tensor1d(tokens.toArray(), "int32").expandDims<tf.Tensor2D>(0),
+    );
+
+    const logits = tf.tidy(() => {
+      const output = this.model.predict(input);
+      if (Array.isArray(output))
+        throw new Error("The model outputs too multiple values");
+      if (output.rank !== 3) throw new Error("The model outputs wrong shape");
+      return output.squeeze<tf.Tensor2D>([0]);
+    });
+    input.dispose();
+
+    const probs = tf.tidy(() =>
+      logits
+        .slice([logits.shape[0] - 1])
+        .squeeze<tf.Tensor1D>([0])
+        .div<tf.Tensor1D>(config.temperature)
+        .softmax(),
+    );
+    logits.dispose();
+
+    const next = tf.tidy(() => {
+      if (config.doSample) {
+        // returns topk biggest values among the `vocab_size` probabilities and the corresponding tokens indices
+        // both shapes are (config.topk,)
+        const { values: topkProbs, indices: topkTokens } = tf.topk(
+          probs,
+          config.topk,
+        );
+        // sample an index from the top-k probabilities
+        // e.g. [[0.1, 0.4, 0.3], [0.1, 0.2, 0.5]] -> [[1], [2]]
+        // note: multinomial does not need the input to sum to 1
+        const selectedIndices = tf.multinomial(
+          topkProbs,
+          1,
+          config.seed,
+          false,
+        ); // (B, )
+        // return the corresponding token from the sampled indices (one per sequence in the batch).
+        // if for some reason the probabilities are NaN, selectedIndices will be out of bounds
+        return topkTokens.gather(selectedIndices).squeeze<tf.Scalar>([0]); // (1)
+      } else {
+        // greedy decoding: return the token with the highest probability
+        return probs.argMax<tf.Scalar>();
+      }
+    });
+    probs.dispose();
+
+    const ret = await next.array();
+    next.dispose();
+    return ret;
+  }
+
+  get config(): Required<GPTConfig> {
+    return this.model.getGPTConfig;
+  }
+  override get weights(): WeightsContainer {
+    return new WeightsContainer(this.model.weights.map((w) => w.read()));
+  }
+
+  override set weights(ws: WeightsContainer) {
+    this.model.setWeights(ws.weights);
+  }
+
+  static deserialize(data: GPTSerialization): Model<"text"> {
+    const model = new GPT(data.config);
+    model.weights = data.weights;
+    return model;
+  }
+
+  serialize(): GPTSerialization {
+    return {
+      weights: this.weights,
+      config: this.config,
+    };
+  }
+  extract(): tf.LayersModel {
+    return this.model;
+  }
+
+  [Symbol.dispose](): void {
+    if (this.model.optimizer !== undefined) {
+      this.model.optimizer.dispose();
+    }
+    const disposeResults = this.model.dispose();
+    if (disposeResults.refCountAfterDispose > 0)
+      debug("model not disposed correctly: %o", disposeResults);
+  }
+}
diff --git a/discojs/src/models/implementations/gpt/index.ts b/discojs/src/models/implementations/gpt/index.ts
index 25de47258..f27b10731 100644
--- a/discojs/src/models/implementations/gpt/index.ts
+++ b/discojs/src/models/implementations/gpt/index.ts
@@ -1,268 +1 @@
-/**
- * Source: https://github.com/zemlyansky/gpt-tfjs and https://github.com/karpathy/build-nanogpt
- * With modifications from @peacefulotter, @lukemovement and the Disco team
- **/
-
-import createDebug from "debug";
-import { List, Range } from "immutable";
-import * as tf from "@tensorflow/tfjs";
-
-import { WeightsContainer } from "#weights/index";
-import { Dataset, Batched } from "#dataset/index";
-import { Model, BatchLogs, EpochLogs } from "#models/index";
-import type { DataFormat } from "#types/index";
-
-import { GPTModel } from "./model.js";
-import evaluate from "./evaluate.js";
-import { DefaultGPTConfig, DefaultGenerationConfig } from "./config.js";
-import type { GPTConfig, GenerationConfig } from "./config.js";
-
-const debug = createDebug("discojs:models:gpt");
-
-export type GPTSerialization = {
-  weights: WeightsContainer;
-  config?: GPTConfig;
-};
-
-export class GPT extends Model<"text"> {
-  readonly datatype = "text" as const;
-
-  private readonly model: GPTModel;
-
-  readonly #contextLength: number;
-  readonly #maxBatchCount: number;
-  readonly #vocabSize: number;
-
-  constructor(
-    partialConfig?: Partial<GPTConfig>,
-    layersModel?: tf.LayersModel,
-  ) {
-    super();
-
-    const model = new GPTModel(partialConfig, layersModel);
-    model.compile();
-    this.model = model;
-
-    this.#contextLength =
-      partialConfig?.contextLength ?? DefaultGPTConfig.contextLength;
-    this.#maxBatchCount = partialConfig?.maxIter ?? DefaultGPTConfig.maxIter;
-    this.#vocabSize = partialConfig?.vocabSize ?? DefaultGPTConfig.vocabSize;
-  }
-
-  /**
-   * The GPT train methods wraps the model.fitDataset call in a for loop to act as a generator (of logs)
-   * This allows for getting logs and stopping training without callbacks.
-   *
-   * @param trainingData training dataset
-   * @param validationData validation dataset
-   * @param epochs the number of passes of the training dataset
-   * @param tracker
-   */
-  override async *train(
-    trainingDataset: Dataset<Batched<DataFormat.ModelEncoded["text"]>>,
-    validationDataset?: Dataset<Batched<DataFormat.ModelEncoded["text"]>>,
-  ): AsyncGenerator<BatchLogs, EpochLogs> {
-    let batchesLogs = List<BatchLogs>();
-    let epochTime = performance.now();
-
-    for await (const [batch, _] of trainingDataset.zip(
-      Range(0, this.#maxBatchCount),
-    )) {
-      const batchLogs = await this.#runBatch(batch);
-
-      yield batchLogs;
-      batchesLogs = batchesLogs.push(batchLogs);
-    }
-
-    const validation =
-      validationDataset && (await this.evaluate(validationDataset));
-    epochTime = performance.now() - epochTime;
-
-    return new EpochLogs(batchesLogs, epochTime, validation);
-  }
-
-  async #runBatch(
-    batch: Batched<DataFormat.ModelEncoded["text"]>,
-  ): Promise<BatchLogs> {
-    const tfBatch = this.#batchToTF(batch);
-
-    let logs: tf.Logs | undefined;
-    await this.model.fitDataset(tf.data.array([tfBatch]), {
-      epochs: 1,
-      verbose: 0, // don't pollute
-      callbacks: {
-        onEpochEnd: (_, cur) => {
-          logs = cur;
-        },
-      },
-    });
-    tf.dispose(tfBatch);
-    if (logs === undefined) throw new Error("batch didn't gave any logs");
-
-    const { loss, acc: accuracy } = logs;
-    if (loss === undefined || isNaN(loss))
-      throw new Error("training loss is undefined or NaN");
-
-    return {
-      accuracy,
-      loss,
-      memoryUsage: tf.memory().numBytes / 1024 / 1024 / 1024,
-    };
-  }
-
-  override async evaluate(
-    dataset: Dataset<Batched<DataFormat.ModelEncoded["text"]>>,
-  ): Promise<Record<"accuracy" | "loss", number>> {
-    const evaluation = await evaluate(
-      this.model,
-      tf.data.generator(
-        async function* (this: GPT) {
-          yield* dataset.map((batch) => this.#batchToTF(batch));
-        }.bind(this),
-      ),
-      this.config.maxEvalBatches,
-    );
-
-    return {
-      accuracy: evaluation.val_acc,
-      loss: evaluation.val_loss,
-    };
-  }
-
-  #batchToTF(batch: Batched<DataFormat.ModelEncoded["text"]>): {
-    xs: tf.Tensor2D;
-    ys: tf.Tensor3D;
-  } {
-    return tf.tidy(() => ({
-      xs: tf.stack(
-        batch.map(([line]) => tf.tensor1d(line.toArray(), "int32")).toArray(),
-      ) as tf.Tensor2D, // cast as stack doesn't type
-      ys: tf.stack(
-        batch
-          .map(([line, next]) =>
-            tf.oneHot(line.shift().push(next).toArray(), this.#vocabSize),
-          )
-          .toArray(),
-      ) as tf.Tensor3D, // cast as oneHot/stack doesn't type
-    }));
-  }
-
-  override async predict(
-    batch: Batched<DataFormat.ModelEncoded["text"][0]>,
-    options?: Partial<GenerationConfig>,
-  ): Promise<Batched<DataFormat.ModelEncoded["text"][1]>> {
-    // overwrite default with user config
-    const config = Object.assign({}, DefaultGenerationConfig, options);
-
-    return List(
-      await Promise.all(
-        batch.map((tokens) => this.#predictSingle(tokens, config)),
-      ),
-    );
-  }
-
-  /**
-   * Generate the next token after the input sequence.
-   * In other words, takes an input tensor of shape (prompt length T) and returns a tensor of shape (T+1)
-   *
-   * @param token input tokens of shape (T,). T is truncated to the model's context length
-   * @param config generation config: temperature, doSample, topk
-   * @returns the next token predicted by the model
-   */
-  async #predictSingle(
-    tokens: DataFormat.ModelEncoded["text"][0],
-    config: GenerationConfig,
-  ): Promise<DataFormat.ModelEncoded["text"][1]> {
-    // slice input tokens if longer than context length
-    tokens = tokens.slice(-this.#contextLength);
-
-    const input = tf.tidy(() =>
-      tf.tensor1d(tokens.toArray(), "int32").expandDims<tf.Tensor2D>(0),
-    );
-
-    const logits = tf.tidy(() => {
-      const output = this.model.predict(input);
-      if (Array.isArray(output))
-        throw new Error("The model outputs too multiple values");
-      if (output.rank !== 3) throw new Error("The model outputs wrong shape");
-      return output.squeeze<tf.Tensor2D>([0]);
-    });
-    input.dispose();
-
-    const probs = tf.tidy(() =>
-      logits
-        .slice([logits.shape[0] - 1])
-        .squeeze<tf.Tensor1D>([0])
-        .div<tf.Tensor1D>(config.temperature)
-        .softmax(),
-    );
-    logits.dispose();
-
-    const next = tf.tidy(() => {
-      if (config.doSample) {
-        // returns topk biggest values among the `vocab_size` probabilities and the corresponding tokens indices
-        // both shapes are (config.topk,)
-        const { values: topkProbs, indices: topkTokens } = tf.topk(
-          probs,
-          config.topk,
-        );
-        // sample an index from the top-k probabilities
-        // e.g. [[0.1, 0.4, 0.3], [0.1, 0.2, 0.5]] -> [[1], [2]]
-        // note: multinomial does not need the input to sum to 1
-        const selectedIndices = tf.multinomial(
-          topkProbs,
-          1,
-          config.seed,
-          false,
-        ); // (B, )
-        // return the corresponding token from the sampled indices (one per sequence in the batch).
-        // if for some reason the probabilities are NaN, selectedIndices will be out of bounds
-        return topkTokens.gather(selectedIndices).squeeze<tf.Scalar>([0]); // (1)
-      } else {
-        // greedy decoding: return the token with the highest probability
-        return probs.argMax<tf.Scalar>();
-      }
-    });
-    probs.dispose();
-
-    const ret = await next.array();
-    next.dispose();
-    return ret;
-  }
-
-  get config(): Required<GPTConfig> {
-    return this.model.getGPTConfig;
-  }
-  override get weights(): WeightsContainer {
-    return new WeightsContainer(this.model.weights.map((w) => w.read()));
-  }
-
-  override set weights(ws: WeightsContainer) {
-    this.model.setWeights(ws.weights);
-  }
-
-  static deserialize(data: GPTSerialization): Model<"text"> {
-    const model = new GPT(data.config);
-    model.weights = data.weights;
-    return model;
-  }
-
-  serialize(): GPTSerialization {
-    return {
-      weights: this.weights,
-      config: this.config,
-    };
-  }
-  extract(): tf.LayersModel {
-    return this.model;
-  }
-
-  [Symbol.dispose](): void {
-    if (this.model.optimizer !== undefined) {
-      this.model.optimizer.dispose();
-    }
-    const disposeResults = this.model.dispose();
-    if (disposeResults.refCountAfterDispose > 0)
-      debug("model not disposed correctly: %o", disposeResults);
-  }
-}
+export { GPT } from "#models/implementations/gpt/gpt";
diff --git a/discojs/src/models/implementations/gpt/layers.spec.ts b/discojs/src/models/implementations/gpt/layers.spec.ts
index b0d6d8a68..2d4141279 100644
--- a/discojs/src/models/implementations/gpt/layers.spec.ts
+++ b/discojs/src/models/implementations/gpt/layers.spec.ts
@@ -1,14 +1,17 @@
 import * as tf from "@tensorflow/tfjs";
 import { afterEach, beforeEach, describe, expect, it } from "vitest";
 
-import type { CausalSelfAttentionConfig, MLPConfig } from "./layers.js";
+import type {
+  CausalSelfAttentionConfig,
+  MLPConfig,
+} from "#models/implementations/gpt/layers";
 import {
   CausalSelfAttention,
   GELU,
   LMEmbedding,
   MLP,
   Range,
-} from "./layers.js";
+} from "#models/implementations/gpt/layers";
 
 describe("GPT Layers", () => {
   // GELU Layer tests
diff --git a/discojs/src/models/implementations/gpt/layers.ts b/discojs/src/models/implementations/gpt/layers.ts
index 83575467c..54c840a53 100644
--- a/discojs/src/models/implementations/gpt/layers.ts
+++ b/discojs/src/models/implementations/gpt/layers.ts
@@ -1,7 +1,6 @@
 import createDebug from "debug";
 import * as tf from "@tensorflow/tfjs";
-import type { GPTConfig } from "./config.js";
-import type { ModelSize } from "./config.js";
+import type { GPTConfig, ModelSize } from "#models/implementations/gpt/config";
 
 const debug = createDebug("discojs:models:gpt:layers");
 
diff --git a/discojs/src/models/implementations/gpt/model.ts b/discojs/src/models/implementations/gpt/model.ts
index c9fd2eb9b..ebee4e92c 100644
--- a/discojs/src/models/implementations/gpt/model.ts
+++ b/discojs/src/models/implementations/gpt/model.ts
@@ -1,11 +1,17 @@
 import createDebug from "debug";
 import * as tf from "@tensorflow/tfjs";
 
-import type { GPTConfig } from "./config.js";
-import { getModelSizes, DefaultGPTConfig } from "./config.js";
-import { getCustomAdam, clipByGlobalNormObj } from "./optimizers.js";
-import evaluate from "./evaluate.js";
-import { GPTArchitecture } from "./layers.js";
+import type { GPTConfig } from "#models/implementations/gpt/config";
+import {
+  getModelSizes,
+  DefaultGPTConfig,
+} from "#models/implementations/gpt/config";
+import {
+  getCustomAdam,
+  clipByGlobalNormObj,
+} from "#models/implementations/gpt/optimizers";
+import evaluate from "#models/implementations/gpt/evaluate";
+import { GPTArchitecture } from "#models/implementations/gpt/layers";
 
 const debug = createDebug("discojs:models:gpt:model");
 
diff --git a/discojs/src/models/implementations/hellaswag.spec.ts b/discojs/src/models/implementations/hellaswag.spec.ts
index e8798f912..05c347a13 100644
--- a/discojs/src/models/implementations/hellaswag.spec.ts
+++ b/discojs/src/models/implementations/hellaswag.spec.ts
@@ -1,7 +1,7 @@
 import { describe, expect, it } from "vitest";
 
-import type { HellaSwagExample } from "./hellaswag.js";
-import { evaluate } from "./hellaswag.js";
+import type { HellaSwagExample } from "#models/implementations/hellaswag";
+import { evaluate } from "#models/implementations/hellaswag";
 import { GPT, Tokenizer } from "#models/index";
 import { ONNXModel } from "#models/onnx";
 
diff --git a/discojs/src/models/implementations/hellaswag.ts b/discojs/src/models/implementations/hellaswag.ts
index b4c643fa3..4e5aa92c3 100644
--- a/discojs/src/models/implementations/hellaswag.ts
+++ b/discojs/src/models/implementations/hellaswag.ts
@@ -1,7 +1,8 @@
 import * as tf from "@tensorflow/tfjs";
 import { List } from "immutable";
-import type { Tokenizer, ONNXModel } from "#models/index";
-import { GPT } from "#models/index";
+import type { Tokenizer } from "#models/tokenizer";
+import type { ONNXModel } from "#models/onnx";
+import { GPT } from "#models/implementations/gpt/index";
 
 export const HELLASWAG_URL =
   "https://raw.githubusercontent.com/rowanz/hellaswag/master/data/hellaswag_val.jsonl";
diff --git a/discojs/src/models/implementations/index.ts b/discojs/src/models/implementations/index.ts
new file mode 100644
index 000000000..982195083
--- /dev/null
+++ b/discojs/src/models/implementations/index.ts
@@ -0,0 +1,4 @@
+export { GPT } from "./gpt/index.js";
+export type { GPTConfig } from "./gpt/config.js";
+export type { HellaSwagDataset, HellaSwagExample } from "./hellaswag.js";
+export { evaluate as evaluate_hellaswag, HELLASWAG_URL } from "./hellaswag.js";
diff --git a/discojs/src/models/implementations/titanicClassifierModel.ts b/discojs/src/models/implementations/titanicClassifierModel.ts
index 9013f1565..4c637a8ef 100644
--- a/discojs/src/models/implementations/titanicClassifierModel.ts
+++ b/discojs/src/models/implementations/titanicClassifierModel.ts
@@ -1,6 +1,6 @@
 import * as tf from "@tensorflow/tfjs";
 
-import { TFJS } from "../index.js";
+import { TFJS } from "#models/tfjs";
 
 export function model() {
   const model = tf.sequential();
diff --git a/discojs/src/models/index.ts b/discojs/src/models/index.ts
index b74b9dc6e..40263a189 100644
--- a/discojs/src/models/index.ts
+++ b/discojs/src/models/index.ts
@@ -3,20 +3,21 @@ export type { BatchLogs, ValidationMetrics } from "./logs.js";
 export { EpochLogs } from "./logs.js";
 export { Tokenizer } from "./tokenizer.js";
 
-export { GPT } from "./implementations/gpt/index.js";
-export { ONNXModel } from "./onnx.js";
-export type { GPTConfig } from "./implementations/gpt/config.js";
 export type {
+  GPTConfig,
   HellaSwagDataset,
   HellaSwagExample,
-} from "./implementations/hellaswag.js";
+} from "./implementations/index.js";
 export {
-  evaluate as evaluate_hellaswag,
+  GPT,
+  evaluate_hellaswag,
   HELLASWAG_URL,
-} from "./implementations/hellaswag.js";
+} from "./implementations/index.js";
+export { ONNXModel } from "./onnx.js";
 export { TFJS } from "./tfjs.js";
 export type { ModelCard } from "./model_card.js";
 export { ModelCardInfo } from "./model_card.js";
 export { fetchModels } from "./model_handler.js";
 
+// eslint-disable-next-line no-restricted-syntax -- namespace re-export acceptable here
 export * as cards from "./cards/index.js";
diff --git a/discojs/src/models/model.ts b/discojs/src/models/model.ts
index 596c6bb5d..b3225f27b 100644
--- a/discojs/src/models/model.ts
+++ b/discojs/src/models/model.ts
@@ -1,8 +1,8 @@
-import { WeightsContainer } from "#weights/index";
-import { Dataset, Batched } from "#dataset/index";
+import type { WeightsContainer } from "#weights/index";
+import type { Dataset, Batched } from "#dataset/index";
 import type { DataFormat, DataType } from "#types/index";
 
-import type { BatchLogs, EpochLogs, ValidationMetrics } from "./logs.js";
+import type { BatchLogs, EpochLogs, ValidationMetrics } from "#models/logs";
 
 /**
  * Trainable predictor
diff --git a/discojs/src/models/model_card.ts b/discojs/src/models/model_card.ts
index 7eeffe3b0..fda563349 100644
--- a/discojs/src/models/model_card.ts
+++ b/discojs/src/models/model_card.ts
@@ -1,6 +1,7 @@
 import { z } from "zod";
-import type { Model } from "#models/index";
-import { DataType, dataTypeValues } from "#types/index";
+import type { Model } from "#models/model";
+import type { DataType } from "#types/index";
+import { dataTypeValues } from "#types/index";
 
 export namespace ModelCardInfo {
   export type ID = string;
diff --git a/discojs/src/models/model_handler.ts b/discojs/src/models/model_handler.ts
index 7b86c9cc3..a629d40f6 100644
--- a/discojs/src/models/model_handler.ts
+++ b/discojs/src/models/model_handler.ts
@@ -3,7 +3,7 @@ import { z } from "zod";
 
 import type { DataType } from "#types/index";
 
-import { ModelCardInfo } from "./model_card.js";
+import { ModelCardInfo } from "#models/model_card";
 
 function urlToModels(base: URL): URL {
   const ret = new URL(base);
diff --git a/discojs/src/models/onnx.spec.ts b/discojs/src/models/onnx.spec.ts
index 688547287..b7eb0018d 100644
--- a/discojs/src/models/onnx.spec.ts
+++ b/discojs/src/models/onnx.spec.ts
@@ -2,8 +2,8 @@ import { describe, expect, it } from "vitest";
 
 import { List } from "immutable";
 import { AutoTokenizer } from "@xenova/transformers";
-import { ONNXModel } from "./onnx.js";
-import { DefaultGenerationConfig } from "./implementations/gpt/config.js";
+import { ONNXModel } from "#models/onnx";
+import { DefaultGenerationConfig } from "#models/implementations/gpt/config";
 
 describe("ONNXModel.predict", { timeout: 50_000 }, () => {
   it("should generate the next token ID from a prompt", async () => {
diff --git a/discojs/src/models/onnx.ts b/discojs/src/models/onnx.ts
index e513f24ca..ae06770ed 100644
--- a/discojs/src/models/onnx.ts
+++ b/discojs/src/models/onnx.ts
@@ -1,17 +1,13 @@
-import type { CausalLMOutput } from "@xenova/transformers";
-import {
-  AutoModelForCausalLM,
-  PreTrainedModel,
-  Tensor,
-} from "@xenova/transformers";
+import type { CausalLMOutput, PreTrainedModel } from "@xenova/transformers";
+import { AutoModelForCausalLM, Tensor } from "@xenova/transformers";
 import { List } from "immutable";
 
 import type { WeightsContainer } from "#weights/index";
 import type { Batched } from "#dataset/index";
 import type { DataFormat } from "#types/index";
-import type { GenerationConfig as TFJSGenerationConfig } from "./implementations/gpt/config.js";
+import type { GenerationConfig as TFJSGenerationConfig } from "#models/implementations/gpt/config";
 import { Model } from "#models/model";
-import { DefaultGenerationConfig } from "./implementations/gpt/config.js";
+import { DefaultGenerationConfig } from "#models/implementations/gpt/config";
 
 export class ONNXModel extends Model<"text"> {
   readonly datatype = "text" as const;
diff --git a/discojs/src/models/tfjs.ts b/discojs/src/models/tfjs.ts
index a31c17e18..d2421b415 100644
--- a/discojs/src/models/tfjs.ts
+++ b/discojs/src/models/tfjs.ts
@@ -1,12 +1,13 @@
 import { List, Map, Range } from "immutable";
 import * as tf from "@tensorflow/tfjs";
 
-import { WeightsContainer } from "#weights/index";
-import { Dataset, Batched } from "#dataset/index";
+import type { Dataset, Batched } from "#dataset/index";
 import type { DataFormat, DataType } from "#types/index";
+import { WeightsContainer } from "#weights/index";
 
-import { Model } from "./model.js";
-import { BatchLogs, EpochLogs } from "./logs.js";
+import type { BatchLogs } from "#models/logs";
+import { Model } from "#models/model";
+import { EpochLogs } from "#models/logs";
 
 type Serialized<D extends DataType> = [D, tf.io.ModelArtifacts];
 
diff --git a/discojs/src/models/tokenizer.spec.ts b/discojs/src/models/tokenizer.spec.ts
index b487323e6..b1e0334e2 100644
--- a/discojs/src/models/tokenizer.spec.ts
+++ b/discojs/src/models/tokenizer.spec.ts
@@ -1,6 +1,6 @@
 import { Repeat } from "immutable";
 import { describe, expect, it } from "vitest";
-import { Tokenizer } from "./tokenizer.js";
+import { Tokenizer } from "#models/tokenizer";
 
 describe("text processing", () => {
   const text = [
diff --git a/discojs/src/privacy.spec.ts b/discojs/src/privacy.spec.ts
index 451a7ea78..9545db314 100644
--- a/discojs/src/privacy.spec.ts
+++ b/discojs/src/privacy.spec.ts
@@ -1,12 +1,12 @@
 import { describe, expect, it } from "vitest";
 
-import { WeightsContainer } from "./index.js";
+import { WeightsContainer } from "#weights/index";
 import {
   frobeniusNorm,
   clipNorm,
   addOptimalNoise,
   getClippingRadius,
-} from "./privacy.js";
+} from "#root/privacy";
 import * as tf from "@tensorflow/tfjs";
 import { List } from "immutable";
 
diff --git a/discojs/src/privacy.ts b/discojs/src/privacy.ts
index c91cb0553..9829b694c 100644
--- a/discojs/src/privacy.ts
+++ b/discojs/src/privacy.ts
@@ -1,8 +1,8 @@
 import * as tf from "@tensorflow/tfjs";
 
-import { WeightsContainer } from "./index.js";
+import { WeightsContainer } from "#weights/index";
 
-import type { WeightNormHistory } from "./training/trainer.js";
+import type { WeightNormHistory } from "#training/types";
 
 /** Computes the Frobenius norm of the given weights. */
 export async function frobeniusNorm(weights: tf.Tensor): Promise<number> {
diff --git a/discojs/src/processing/image.spec.ts b/discojs/src/processing/image.spec.ts
index ebe9af8ce..2d8e9bb05 100644
--- a/discojs/src/processing/image.spec.ts
+++ b/discojs/src/processing/image.spec.ts
@@ -1,7 +1,9 @@
 import { Repeat, Seq } from "immutable";
 import { describe, expect, it } from "vitest";
-import { Image } from "../index.js";
-import { removeAlpha, resize } from "./image.js";
+
+import { Image } from "#dataset/index";
+
+import { removeAlpha, resize } from "#processing/image";
 
 describe("resize", () => {
   it("doesn't change with same image dimensions", () => {
diff --git a/discojs/src/processing/image.ts b/discojs/src/processing/image.ts
index 4a78079d1..cc45caff3 100644
--- a/discojs/src/processing/image.ts
+++ b/discojs/src/processing/image.ts
@@ -2,7 +2,7 @@ import { Repeat, Seq } from "immutable";
 import { createJimp } from "@jimp/core";
 import * as jimpResize from "@jimp/plugin-resize";
 
-import { Image } from "../index.js";
+import { Image } from "#dataset/index";
 
 /** Image where intensity is represented in the range 0..1 */
 export class NormalizedImage<
@@ -90,15 +90,15 @@ export function removeAlpha<W extends number, H extends number>(
 }
 
 /** Convert monochrome images to multicolor */
-export function expandToMulticolor<W extends number, H extends number>(
+function expandToMulticolor<W extends number, H extends number>(
   image: Image<1, W, H>,
 ): Image<3, W, H>;
-export function expandToMulticolor<
+function expandToMulticolor<
   D extends 3 | 4,
   W extends number,
   H extends number,
 >(image: Image<1 | D, W, H>): Image<D, W, H>;
-export function expandToMulticolor<W extends number, H extends number>(
+function expandToMulticolor<W extends number, H extends number>(
   image: Image<1 | 3 | 4, W, H>,
 ): Image<3 | 4, W, H> {
   switch (image.depth) {
diff --git a/discojs/src/processing/index.ts b/discojs/src/processing/index.ts
index 751017ee4..ea0fe080c 100644
--- a/discojs/src/processing/index.ts
+++ b/discojs/src/processing/index.ts
@@ -1,136 +1,7 @@
 /** Dataset shapers, convenient to map with */
-
-import { List } from "immutable";
-
-import type {
-  Dataset,
-  DataFormat,
-  DataType,
-  Tabular,
-  Task,
-  Network,
-} from "../index.js";
-
-import * as processing from "./index.js";
-
-export * from "./image.js";
-export * from "./tabular.js";
-
-export function preprocess<D extends DataType, N extends Network>(
-  task: Task<D, N>,
-  dataset: Dataset<DataFormat.Raw[D]>,
-): Dataset<DataFormat.ModelEncoded[D]> {
-  switch (task.dataType) {
-    case "image": {
-      // cast as typescript doesn't reduce generic type
-      const d = dataset as Dataset<DataFormat.Raw["image"]>;
-      const { IMAGE_H, IMAGE_W, LABEL_LIST } = task.trainingInformation;
-
-      return d.map(([image, label]) => [
-        processing.normalize(
-          processing.removeAlpha(processing.resize(IMAGE_W, IMAGE_H, image)),
-        ),
-        processing.indexInList(label, LABEL_LIST),
-      ]) as Dataset<DataFormat.ModelEncoded[D]>;
-    }
-    case "tabular": {
-      // cast as typescript doesn't reduce generic type
-      const d = dataset as Dataset<DataFormat.Raw["tabular"]>;
-      const { inputColumns, outputColumn } = task.trainingInformation;
-
-      return d.map((row) => {
-        const output = processing.extractColumn(row, outputColumn);
-
-        return [
-          extractToNumbers(inputColumns, row),
-          // TODO sanitization doesn't care about column distribution
-          output !== "" ? processing.convertToNumber(output) : 0,
-        ];
-      }) as Dataset<DataFormat.ModelEncoded[D]>;
-    }
-    case "text": {
-      // cast as typescript doesn't reduce generic type
-      const d = dataset as Dataset<DataFormat.Raw["text"]>;
-
-      const { contextLength, tokenizer } = task.trainingInformation;
-
-      return d
-        .map((text) => tokenizer.tokenize(text))
-        .flatten()
-        .batch(contextLength + 1, 1)
-        .map((tokens) => [tokens.pop(), tokens.last()]) as Dataset<
-        DataFormat.ModelEncoded[D]
-      >;
-    }
-  }
-}
-
-export function preprocessWithoutLabel<D extends DataType>(
-  task: Task<D, Network>,
-  dataset: Dataset<DataFormat.RawWithoutLabel[D]>,
-): Dataset<DataFormat.ModelEncoded[D][0]> {
-  switch (task.dataType) {
-    case "image": {
-      // cast as typescript doesn't reduce generic type
-      const d = dataset as Dataset<DataFormat.RawWithoutLabel["image"]>;
-      const { IMAGE_H, IMAGE_W } = task.trainingInformation;
-
-      return d.map((image) =>
-        processing.normalize(
-          processing.removeAlpha(processing.resize(IMAGE_W, IMAGE_H, image)),
-        ),
-      );
-    }
-    case "tabular": {
-      // cast as typescript doesn't reduce generic type
-      const d = dataset as Dataset<DataFormat.Raw["tabular"]>;
-      const { inputColumns } = task.trainingInformation;
-
-      return d.map((row) => extractToNumbers(inputColumns, row));
-    }
-    case "text": {
-      // cast as typescript doesn't reduce generic type
-      const d = dataset as Dataset<DataFormat.Raw["text"]>;
-
-      const { contextLength, tokenizer } = task.trainingInformation;
-
-      return d
-        .map((text) => tokenizer.tokenize(text))
-        .flatten()
-        .batch(contextLength);
-    }
-  }
-}
-
-export function postprocess<D extends DataType>(
-  task: Task<D, Network>,
-  encoded: DataFormat.ModelEncoded[D][1],
-): DataFormat.Inferred[D] {
-  switch (task.dataType) {
-    case "image": {
-      const labels = List(task.trainingInformation.LABEL_LIST);
-
-      const v = labels.get(encoded);
-      if (v === undefined) throw new Error("index not found in labels");
-      return v as DataFormat.Inferred[D];
-    }
-    case "tabular": {
-      return encoded as DataFormat.Inferred[D];
-    }
-    case "text": {
-      return task.trainingInformation.tokenizer.decode([
-        encoded,
-      ]) as DataFormat.Inferred[D];
-    }
-  }
-}
-
-function extractToNumbers(columns: Iterable<string>, row: Tabular) {
-  return (
-    List(columns)
-      .map((column) => processing.extractColumn(row, column))
-      // TODO sanitization doesn't care about column distribution
-      .map((v) => (v !== "" ? v : "0"))
-      .map(processing.convertToNumber)
-  );
-}
+export { extractColumn } from "./tabular.js";
+export {
+  preprocess,
+  preprocessWithoutLabel,
+  postprocess,
+} from "./processing.js";
diff --git a/discojs/src/processing/index.spec.ts b/discojs/src/processing/processing.spec.ts
similarity index 87%
rename from discojs/src/processing/index.spec.ts
rename to discojs/src/processing/processing.spec.ts
index 19aa2b47a..bf995b99b 100644
--- a/discojs/src/processing/index.spec.ts
+++ b/discojs/src/processing/processing.spec.ts
@@ -1,9 +1,9 @@
 import { describe, expect, it } from "vitest";
 
-import type { Task } from "../index.js";
-import { Dataset } from "../index.js";
+import { preprocess } from "#processing/processing";
 
-import { preprocess } from "./index.js";
+import type { Task } from "#task/index";
+import { Dataset } from "#dataset/index";
 
 describe("preprocess", () => {
   it("throws on missing column in tabular", async () => {
diff --git a/discojs/src/processing/processing.ts b/discojs/src/processing/processing.ts
new file mode 100644
index 000000000..3cc9e981c
--- /dev/null
+++ b/discojs/src/processing/processing.ts
@@ -0,0 +1,127 @@
+import { List } from "immutable";
+
+import type { Task } from "#task/index";
+import type { Dataset, Tabular } from "#dataset/index";
+import type { DataType, DataFormat, Network } from "#types/index";
+
+import { normalize, removeAlpha, resize } from "#processing/image";
+import {
+  indexInList,
+  extractColumn,
+  convertToNumber,
+} from "#processing/tabular";
+
+export function preprocess<D extends DataType, N extends Network>(
+  task: Task<D, N>,
+  dataset: Dataset<DataFormat.Raw[D]>,
+): Dataset<DataFormat.ModelEncoded[D]> {
+  switch (task.dataType) {
+    case "image": {
+      // cast as typescript doesn't reduce generic type
+      const d = dataset as Dataset<DataFormat.Raw["image"]>;
+      const { IMAGE_H, IMAGE_W, LABEL_LIST } = task.trainingInformation;
+
+      return d.map(([image, label]) => [
+        normalize(removeAlpha(resize(IMAGE_W, IMAGE_H, image))),
+        indexInList(label, LABEL_LIST),
+      ]) as Dataset<DataFormat.ModelEncoded[D]>;
+    }
+    case "tabular": {
+      // cast as typescript doesn't reduce generic type
+      const d = dataset as Dataset<DataFormat.Raw["tabular"]>;
+      const { inputColumns, outputColumn } = task.trainingInformation;
+
+      return d.map((row) => {
+        const output = extractColumn(row, outputColumn);
+
+        return [
+          extractToNumbers(inputColumns, row),
+          // TODO sanitization doesn't care about column distribution
+          output !== "" ? convertToNumber(output) : 0,
+        ];
+      }) as Dataset<DataFormat.ModelEncoded[D]>;
+    }
+    case "text": {
+      // cast as typescript doesn't reduce generic type
+      const d = dataset as Dataset<DataFormat.Raw["text"]>;
+
+      const { contextLength, tokenizer } = task.trainingInformation;
+
+      return d
+        .map((text) => tokenizer.tokenize(text))
+        .flatten()
+        .batch(contextLength + 1, 1)
+        .map((tokens) => [tokens.pop(), tokens.last()]) as Dataset<
+        DataFormat.ModelEncoded[D]
+      >;
+    }
+  }
+}
+
+export function preprocessWithoutLabel<D extends DataType>(
+  task: Task<D, Network>,
+  dataset: Dataset<DataFormat.RawWithoutLabel[D]>,
+): Dataset<DataFormat.ModelEncoded[D][0]> {
+  switch (task.dataType) {
+    case "image": {
+      // cast as typescript doesn't reduce generic type
+      const d = dataset as Dataset<DataFormat.RawWithoutLabel["image"]>;
+      const { IMAGE_H, IMAGE_W } = task.trainingInformation;
+
+      return d.map((image) =>
+        normalize(removeAlpha(resize(IMAGE_W, IMAGE_H, image))),
+      );
+    }
+    case "tabular": {
+      // cast as typescript doesn't reduce generic type
+      const d = dataset as Dataset<DataFormat.Raw["tabular"]>;
+      const { inputColumns } = task.trainingInformation;
+
+      return d.map((row) => extractToNumbers(inputColumns, row));
+    }
+    case "text": {
+      // cast as typescript doesn't reduce generic type
+      const d = dataset as Dataset<DataFormat.Raw["text"]>;
+
+      const { contextLength, tokenizer } = task.trainingInformation;
+
+      return d
+        .map((text) => tokenizer.tokenize(text))
+        .flatten()
+        .batch(contextLength);
+    }
+  }
+}
+
+export function postprocess<D extends DataType>(
+  task: Task<D, Network>,
+  encoded: DataFormat.ModelEncoded[D][1],
+): DataFormat.Inferred[D] {
+  switch (task.dataType) {
+    case "image": {
+      const labels = List(task.trainingInformation.LABEL_LIST);
+
+      const v = labels.get(encoded);
+      if (v === undefined) throw new Error("index not found in labels");
+      return v as DataFormat.Inferred[D];
+    }
+    case "tabular": {
+      return encoded as DataFormat.Inferred[D];
+    }
+    case "text": {
+      return task.trainingInformation.tokenizer.decode([
+        encoded,
+      ]) as DataFormat.Inferred[D];
+    }
+  }
+}
+
+function extractToNumbers(columns: Iterable<string>, row: Tabular) {
+  return (
+    List(columns)
+      .map((column) => extractColumn(row, column))
+      // TODO sanitization doesn't care about column distribution
+      .map((v) => (v !== "" ? v : "0"))
+      .map(convertToNumber)
+  );
+}
diff --git a/discojs/src/processing/tabular.ts b/discojs/src/processing/tabular.ts
index 685baca03..9f9dcc58a 100644
--- a/discojs/src/processing/tabular.ts
+++ b/discojs/src/processing/tabular.ts
@@ -1,4 +1,4 @@
-import { List } from "immutable";
+import type { List } from "immutable";
 
 /**
  * Convert a string to a number
diff --git a/discojs/src/serialization/index.ts b/discojs/src/serialization/index.ts
index 1dccea6e6..f103ac5c5 100644
--- a/discojs/src/serialization/index.ts
+++ b/discojs/src/serialization/index.ts
@@ -1,15 +1,11 @@
-export * as model from "./model.js";
-export * as task from "./task.js";
-export * as weights from "./weights.js";
+export { encode as modelEncode, decode as modelDecode } from "./model.js";
+export {
+  serializeToJSON as serializeTaskToJSON,
+  deserializeFromJSON as deserializeTaskFromJSON,
+} from "./task.js";
+export { encode as weightsEncode, decode as weightsDecode } from "./weights.js";
 
 export type { Encoded } from "./coder.js";
 export { isEncoded } from "./coder.js";
 
-export type JSON =
-  | null
-  | undefined
-  | boolean
-  | number
-  | string
-  | JSON[]
-  | { [_: string]: JSON };
+export type { JSONLike } from "./json_like.js";
diff --git a/discojs/src/serialization/json_like.ts b/discojs/src/serialization/json_like.ts
new file mode 100644
index 000000000..30505b626
--- /dev/null
+++ b/discojs/src/serialization/json_like.ts
@@ -0,0 +1,8 @@
+export type JSONLike =
+  | null
+  | undefined
+  | boolean
+  | number
+  | string
+  | JSONLike[]
+  | { [_: string]: JSONLike };
diff --git a/discojs/src/serialization/model.spec.ts b/discojs/src/serialization/model.spec.ts
index 6cf0e3bf8..a11f21d5f 100644
--- a/discojs/src/serialization/model.spec.ts
+++ b/discojs/src/serialization/model.spec.ts
@@ -1,9 +1,12 @@
 import * as tf from "@tensorflow/tfjs";
 import { assert, describe, expect, it } from "vitest";
 
-import type { DataType, Model } from "../index.js";
-import { models, serialization } from "../index.js";
-import type { GPTConfig } from "../models/index.js";
+import type { DataType } from "#types/index";
+import type { Model, GPTConfig } from "#models/index";
+import { GPT, TFJS } from "#models/index";
+
+import { encode, decode } from "#serialization/model";
+import { isEncoded } from "#serialization/coder";
 
 async function getRawWeights(
   model: Model<DataType>,
@@ -30,16 +33,14 @@ describe("serialization", () => {
       ],
     });
     rawModel.compile({ optimizer: "sgd", loss: "hinge" });
-    const model = new models.TFJS("image", rawModel);
+    const model = new TFJS("image", rawModel);
 
-    const encoded = await serialization.model.encode(model);
-    assert.isTrue(serialization.isEncoded(encoded));
-    const decoded = await serialization.model.decode(encoded);
+    const encoded = await encode(model);
+    assert.isTrue(isEncoded(encoded));
+    const decoded = await decode(encoded);
 
-    expect(decoded).to.be.an.instanceof(models.TFJS);
-    expect((decoded as models.TFJS<"image" | "tabular">).datatype).to.equal(
-      "image",
-    );
+    expect(decoded).to.be.an.instanceof(TFJS);
+    expect((decoded as TFJS<"image" | "tabular">).datatype).to.equal("image");
     assert.sameDeepOrderedMembers(
       await getRawWeights(model),
       await getRawWeights(decoded),
@@ -55,13 +56,13 @@ describe("serialization", () => {
       maxEvalBatches: 10,
       contextLength: 8,
     };
-    const model = new models.GPT(config);
+    const model = new GPT(config);
 
-    const encoded = await serialization.model.encode(model);
-    assert.isTrue(serialization.isEncoded(encoded));
-    const decoded = await serialization.model.decode(encoded);
+    const encoded = await encode(model);
+    assert.isTrue(isEncoded(encoded));
+    const decoded = await decode(encoded);
 
-    assert.instanceOf(decoded, models.GPT);
+    assert.instanceOf(decoded, GPT);
 
     assert.sameDeepOrderedMembers(
       await getRawWeights(model),
diff --git a/discojs/src/serialization/model.ts b/discojs/src/serialization/model.ts
index 5552bfca1..8a56e524c 100644
--- a/discojs/src/serialization/model.ts
+++ b/discojs/src/serialization/model.ts
@@ -1,11 +1,16 @@
 import type tf from "@tensorflow/tfjs";
 
-import type { DataType, Model } from "../index.js";
-import { models, serialization } from "../index.js";
-import { GPTConfig } from "../models/index.js";
+import { encode as w_encode, decode as w_decode } from "#serialization/weights";
+import { GPT, TFJS } from "#models/index";
+import type { Model, GPTConfig } from "#models/index";
+import type { DataType } from "#types/index";
 
-import * as coder from "./coder.js";
-import { Encoded, isEncoded } from "./coder.js";
+import type { Encoded } from "#serialization/coder";
+import {
+  encode as encodeGeneric,
+  decode as decodeGeneric,
+  isEncoded,
+} from "#serialization/coder";
 
 const Type = {
   TFJS: 0,
@@ -14,14 +19,14 @@ const Type = {
 
 export async function encode(model: Model<DataType>): Promise<Encoded> {
   switch (true) {
-    case model instanceof models.TFJS: {
+    case model instanceof TFJS: {
       const serialized = await model.serialize();
-      return coder.encode([Type.TFJS, ...serialized]);
+      return encodeGeneric([Type.TFJS, ...serialized]);
     }
-    case model instanceof models.GPT: {
+    case model instanceof GPT: {
       const { weights, config } = model.serialize();
-      const serializedWeights = await serialization.weights.encode(weights);
-      return coder.encode([Type.GPT, serializedWeights, config]);
+      const serializedWeights = await w_encode(weights);
+      return encodeGeneric([Type.GPT, serializedWeights, config]);
     }
     default:
       throw new Error("unknown model type");
@@ -29,7 +34,7 @@ export async function encode(model: Model<DataType>): Promise<Encoded> {
 }
 
 export async function decode(encoded: Encoded): Promise<Model<DataType>> {
-  const raw = coder.decode(encoded);
+  const raw = decodeGeneric(encoded);
 
   if (!Array.isArray(raw) || raw.length < 2) {
     throw new Error(
@@ -61,7 +66,7 @@ export async function decode(encoded: Encoded): Promise<Model<DataType>> {
           throw new Error("invalid TFJS model encoding: invalid DataType");
       }
 
-      return await models.TFJS.deserialize([
+      return await TFJS.deserialize([
         datatype,
         // TODO totally unsafe casting
         rawModel as tf.io.ModelArtifacts,
@@ -83,8 +88,8 @@ export async function decode(encoded: Encoded): Promise<Model<DataType>> {
         throw new Error(
           "invalid encoding, gpt-tfjs model weights should be an encoding of its weights",
         );
-      const weights = serialization.weights.decode(rawModel);
-      return models.GPT.deserialize({ weights, config });
+      const weights = w_decode(rawModel);
+      return GPT.deserialize({ weights, config });
     }
     default:
       throw new Error("invalid encoding, model type unrecognized");
diff --git a/discojs/src/serialization/task.spec.ts b/discojs/src/serialization/task.spec.ts
index 2c3061f0a..9c5cfabcf 100644
--- a/discojs/src/serialization/task.spec.ts
+++ b/discojs/src/serialization/task.spec.ts
@@ -1,12 +1,13 @@
 import { expect, it } from "vitest";
 
-import { serialization, defaultTasks } from "../index.js";
+import { defaultTasks } from "#root/index";
+import { deserializeFromJSON, serializeToJSON } from "#serialization/task";
 
 it("can encode what it decodes", async () => {
   const task = await defaultTasks.wikitext.getTask();
 
-  const serialized = serialization.task.serializeToJSON(task);
-  const deserialized = await serialization.task.deserializeFromJSON(serialized);
+  const serialized = serializeToJSON(task);
+  const deserialized = await deserializeFromJSON(serialized);
 
   expect(deserialized).to.be.deep.equal(task);
 });
diff --git a/discojs/src/serialization/task.ts b/discojs/src/serialization/task.ts
index ea1a36f47..6d292da1c 100644
--- a/discojs/src/serialization/task.ts
+++ b/discojs/src/serialization/task.ts
@@ -1,10 +1,11 @@
 import { z } from "zod";
-import type { DataType, Network } from "../index.js";
-import { Task, Tokenizer } from "../index.js";
+import type { DataType, Network } from "#types/index";
+import { Task } from "#task/task";
+import { Tokenizer } from "#models/index";
 
-import type { JSON } from "./index.js";
+import type { JSONLike } from "#serialization/json_like";
 
-export function serializeToJSON(task: Task<DataType, Network>): JSON {
+export function serializeToJSON(task: Task<DataType, Network>): JSONLike {
   switch (task.dataType) {
     case "image":
     case "tabular":
@@ -23,7 +24,7 @@ export function serializeToJSON(task: Task<DataType, Network>): JSON {
 
 // Throws if an error serialized object is malformed
 export async function deserializeFromJSON(
-  serialized: JSON,
+  serialized: JSONLike,
 ): Promise<Task<DataType, Network>> {
   return await z
     .looseObject({
diff --git a/discojs/src/serialization/weights.spec.ts b/discojs/src/serialization/weights.spec.ts
index b13aea6ad..f3a0bc998 100644
--- a/discojs/src/serialization/weights.spec.ts
+++ b/discojs/src/serialization/weights.spec.ts
@@ -1,14 +1,16 @@
 import { assert, describe, it } from "vitest";
 
-import { WeightsContainer, serialization } from "../index.js";
+import { WeightsContainer } from "#weights/index";
+import { isEncoded } from "#serialization/coder";
+import { encode, decode } from "#serialization/weights";
 
 describe("weights", () => {
   it("can encode what it decodes", async () => {
     const weights = WeightsContainer.of([1], [2], [3]);
 
-    const encoded = await serialization.weights.encode(weights);
-    assert.isTrue(serialization.isEncoded(encoded));
-    const decoded = serialization.weights.decode(encoded);
+    const encoded = await encode(weights);
+    assert.isTrue(isEncoded(encoded));
+    const decoded = decode(encoded);
 
     assert.sameDeepOrderedMembers(
       Array.from(
diff --git a/discojs/src/serialization/weights.ts b/discojs/src/serialization/weights.ts
index 52af2c8a4..a0441defc 100644
--- a/discojs/src/serialization/weights.ts
+++ b/discojs/src/serialization/weights.ts
@@ -1,9 +1,12 @@
 import * as tf from "@tensorflow/tfjs";
 
-import { WeightsContainer } from "../index.js";
+import { WeightsContainer } from "#weights/index";
 
-import { Encoded } from "./coder.js";
-import * as coder from "./coder.js";
+import type { Encoded } from "#serialization/coder";
+import {
+  encode as encodeGeneric,
+  decode as decodeGeneric,
+} from "#serialization/coder";
 
 type Serialized = {
   shape: number[];
@@ -34,11 +37,11 @@ export async function encode(weights: WeightsContainer): Promise<Encoded> {
     })),
   );
 
-  return coder.encode(serialized);
+  return encodeGeneric(serialized);
 }
 
 export function decode(encoded: Encoded): WeightsContainer {
-  const raw = coder.decode(encoded);
+  const raw = decodeGeneric(encoded);
 
   if (!(Array.isArray(raw) && raw.every(isSerialized)))
     throw new Error("expected to decode an array of serialized weights");
diff --git a/discojs/src/task/display_information.ts b/discojs/src/task/display_information.ts
index b03b91da2..c87c70a26 100644
--- a/discojs/src/task/display_information.ts
+++ b/discojs/src/task/display_information.ts
@@ -1,6 +1,6 @@
 import { z } from "zod";
 
-import type { DataType } from "../types/index.js";
+import type { DataType } from "#types/index";
 
 export namespace DisplayInformation {
   export const baseSchema = z.object({
diff --git a/discojs/src/task/index.ts b/discojs/src/task/index.ts
index 4073e158c..51e6d69cf 100644
--- a/discojs/src/task/index.ts
+++ b/discojs/src/task/index.ts
@@ -1,6 +1,5 @@
 export { Task } from "./task.js";
 export type { TaskProvider } from "./task_provider.js";
-export { DisplayInformation } from "./display_information.js";
 export { TrainingInformation } from "./training_information.js";
 
 export { pushTask, fetchTasks } from "./task_handler.js";
diff --git a/discojs/src/task/task.ts b/discojs/src/task/task.ts
index 2903239cc..12ab11fb1 100644
--- a/discojs/src/task/task.ts
+++ b/discojs/src/task/task.ts
@@ -1,9 +1,9 @@
 import { z } from "zod";
 
-import type { DataType, Network } from "../index.js";
+import type { DataType, Network } from "#types/index";
 
-import { DisplayInformation } from "./display_information.js";
-import { TrainingInformation } from "./training_information.js";
+import { DisplayInformation } from "#task/display_information";
+import { TrainingInformation } from "#task/training_information";
 
 export namespace Task {
   export type ID = string;
diff --git a/discojs/src/task/task_handler.ts b/discojs/src/task/task_handler.ts
index 1fb6273ce..bcb4d1f00 100644
--- a/discojs/src/task/task_handler.ts
+++ b/discojs/src/task/task_handler.ts
@@ -1,10 +1,15 @@
 import { Map, Seq } from "immutable";
 
-import type { DataType, Model, Network } from "../index.js";
-import { serialization } from "../index.js";
-import type { ModelCardInfo } from "#models/model_card";
+import type { DataType, Network } from "#types/index";
+import type { Model, ModelCardInfo } from "#models/index";
+import type { JSONLike } from "#serialization/index";
+import {
+  serializeTaskToJSON,
+  deserializeTaskFromJSON,
+  modelEncode,
+} from "#serialization/index";
 
-import type { Task } from "./task.js";
+import type { Task } from "#task/task";
 
 function urlToTasks(base: URL): URL {
   const ret = new URL(base);
@@ -27,11 +32,9 @@ export async function pushTask<D extends DataType>(
     method: "POST",
     headers: { "Content-Type": "application/json" },
     body: JSON.stringify({
-      task: serialization.task.serializeToJSON(task),
+      task: serializeTaskToJSON(task),
       model:
-        typeof model === "string"
-          ? model
-          : [...(await serialization.model.encode(model))],
+        typeof model === "string" ? model : [...(await modelEncode(model))],
     }),
   });
   if (!response.ok) throw new Error(`fetch: HTTP status ${response.status}`);
@@ -42,7 +45,7 @@ export async function fetchTasks(
 ): Promise<Map<Task.ID, Task<DataType, Network>>> {
   const response = await fetch(urlToTasks(base));
   if (!response.ok) throw new Error(`fetch: HTTP status ${response.status}`);
-  const json = (await response.json()) as serialization.JSON;
+  const json = (await response.json()) as JSONLike;
 
   if (!Array.isArray(json))
     throw new Error("invalid tasks response: expected a JSON array");
@@ -50,11 +53,9 @@ export async function fetchTasks(
 
   try {
     return Map(
-      Seq(
-        await Promise.all(
-          arr.map((t) => serialization.task.deserializeFromJSON(t)),
-        ),
-      ).map((t) => [t.id, t]),
+      Seq(await Promise.all(arr.map((t) => deserializeTaskFromJSON(t)))).map(
+        (t) => [t.id, t],
+      ),
     );
   } catch (cause) {
     throw new Error("invalid tasks response: unable to parse all tasks", {
diff --git a/discojs/src/task/task_provider.ts b/discojs/src/task/task_provider.ts
index 666856f85..4609a9746 100644
--- a/discojs/src/task/task_provider.ts
+++ b/discojs/src/task/task_provider.ts
@@ -1,4 +1,6 @@
-import type { DataType, Network, Task, ModelCard } from "../index.js";
+import type { DataType, Network } from "#types/index";
+import type { ModelCard } from "#models/index";
+import type { Task } from "#task/task";
 
 export interface TaskProvider<D extends DataType, N extends Network> {
   getTask(): Promise<Task<D, N>>;
diff --git a/discojs/src/task/training_information.ts b/discojs/src/task/training_information.ts
index a8f041ebb..be7e4bd01 100644
--- a/discojs/src/task/training_information.ts
+++ b/discojs/src/task/training_information.ts
@@ -1,7 +1,7 @@
 import { z } from "zod";
 
-import type { DataType, Network } from "../index.js";
-import { Tokenizer } from "../index.js";
+import type { DataType, Network } from "#types/index";
+import { Tokenizer } from "#models/index";
 
 const privacySchema = z.object({
   // reduce training accuracy and improve privacy.
diff --git a/discojs/src/training/disco.ts b/discojs/src/training/disco.ts
index 4a124f9bf..0abd4a38e 100644
--- a/discojs/src/training/disco.ts
+++ b/discojs/src/training/disco.ts
@@ -1,27 +1,24 @@
-import {
-  async_iterator,
-  client as clients,
-  BatchLogs,
-  ConsoleLogger,
-  EpochLogs,
-  Logger,
-  processing,
-  Dataset,
-} from "../index.js";
-import type {
-  Batched,
-  DataFormat,
-  DataType,
-  Model,
-  Network,
-  Task,
-} from "../index.js";
-import type { Aggregator } from "../aggregator/index.js";
-import { getAggregator } from "../aggregator/index.js";
-import { enumerate, split } from "../utils/async_iterator.js";
-import { EventEmitter } from "../utils/event_emitter.js";
-
-import { RoundLogs, Trainer } from "./trainer.js";
+import type { Model } from "#models/index";
+import type { DataType, DataFormat, Network } from "#types/index";
+import type { Task } from "#task/index";
+import type { Batched } from "#dataset/index";
+import type { Aggregator } from "#aggregator/index";
+
+import { Dataset } from "#dataset/index";
+import type { Logger } from "#logging/index";
+import { ConsoleLogger } from "#logging/index";
+import type { BatchLogs, EpochLogs } from "#models/index";
+import { getAggregator } from "#aggregator/index";
+import { enumerate, split } from "#utils/async_iterator";
+import { EventEmitter } from "#utils/event_emitter";
+
+import * as clients from "#client/index";
+import * as processing from "#processing/index";
+import * as async_iterator from "#utils/async_iterator";
+
+import type { RoundLogs } from "#training/trainer";
+import { Trainer } from "#training/trainer";
+import type { RoundStatus, SummaryLogs } from "#training/types";
 
 interface DiscoConfig<N extends Network> {
   scheme: N;
@@ -37,25 +34,6 @@ interface DiscoConfig<N extends Network> {
   preprocessOnce: boolean;
 }
 
-export type SummaryLogs = {
-  round: number;
-  epoch: number;
-  trainingLoss: number;
-  trainingAccuracy: number;
-  peakMemory: number;
-  epochTime: number;
-  roundValidationLoss?: number;
-  roundValidationAccuracy?: number;
-  validationLoss?: number;
-  validationAccuracy?: number;
-};
-
-export type RoundStatus =
-  | "not enough participants" // Server notification to wait for more participants
-  | "updating model" // fetching/aggregating local updates into a global model
-  | "local training" // Training the model locally
-  | "connecting to peers"; // for decentralized only, fetch the server's list of participating peers
-
 function buildSummaryLog(
   roundNum: number,
   epochNum: number,
diff --git a/discojs/src/training/index.ts b/discojs/src/training/index.ts
index 1c10a3fef..0cad8570a 100644
--- a/discojs/src/training/index.ts
+++ b/discojs/src/training/index.ts
@@ -1,2 +1,3 @@
-export { Disco, RoundStatus, SummaryLogs } from "./disco.js";
-export { RoundLogs, Trainer } from "./trainer.js";
+export { Disco } from "./disco.js";
+export type { RoundStatus, SummaryLogs } from "./types.js";
+export type { RoundLogs } from "./trainer.js";
diff --git a/discojs/src/training/trainer.ts b/discojs/src/training/trainer.ts
index 44c8f8cd3..023b03766 100644
--- a/discojs/src/training/trainer.ts
+++ b/discojs/src/training/trainer.ts
@@ -1,22 +1,20 @@
 import * as tf from "@tensorflow/tfjs";
 import { List, Repeat } from "immutable";
 
-import {
-  Batched,
-  BatchLogs,
-  Dataset,
-  DataFormat,
-  DataType,
-  EpochLogs,
-  Model,
-  Task,
-  WeightsContainer,
-  Network,
-  ValidationMetrics,
-} from "../index.js";
-import { privacy } from "../index.js";
-import { Client } from "../client/index.js";
-import * as async_iterator from "../utils/async_iterator.js";
+import type { Model } from "#models/index";
+import type { DataFormat, DataType, Network } from "#types/index";
+import type { Batched } from "#dataset/index";
+import type { Task } from "#task/index";
+
+import type { Dataset } from "#dataset/index";
+import type { BatchLogs, EpochLogs, ValidationMetrics } from "#models/index";
+import { WeightsContainer } from "#weights/index";
+import type { Client } from "#client/index";
+
+import * as async_iterator from "#utils/async_iterator";
+import * as privacy from "#root/privacy";
+
+import type { WeightNormHistory } from "#training/types";
 
 export interface RoundLogs {
   epochs: List<EpochLogs>;
@@ -25,8 +23,6 @@ export interface RoundLogs {
 }
 
 /** List of weight update norms */
-export type WeightNormHistory = List<List<number>>;
-
 function appendWeightHistory(
   weightNormHistory: WeightNormHistory,
   wc: number[],
diff --git a/discojs/src/training/types.ts b/discojs/src/training/types.ts
new file mode 100644
index 000000000..87be9e6ad
--- /dev/null
+++ b/discojs/src/training/types.ts
@@ -0,0 +1,22 @@
+import type { List } from "immutable";
+
+export type WeightNormHistory = List<List<number>>;
+
+export type SummaryLogs = {
+  round: number;
+  epoch: number;
+  trainingLoss: number;
+  trainingAccuracy: number;
+  peakMemory: number;
+  epochTime: number;
+  roundValidationLoss?: number;
+  roundValidationAccuracy?: number;
+  validationLoss?: number;
+  validationAccuracy?: number;
+};
+
+export type RoundStatus =
+  | "not enough participants" // Server notification to wait for more participants
+  | "updating model" // fetching/aggregating local updates into a global model
+  | "local training" // Training the model locally
+  | "connecting to peers"; // for decentralized only, fetch the server's list of participating peers
diff --git a/discojs/src/types/data_format.ts b/discojs/src/types/data_format.ts
index 7047d4be8..6cce455aa 100644
--- a/discojs/src/types/data_format.ts
+++ b/discojs/src/types/data_format.ts
@@ -1,12 +1,7 @@
-import { List } from "immutable";
+import type { List } from "immutable";
 
-import type {
-  Image,
-  processing,
-  Tabular,
-  Text,
-  TokenizedText,
-} from "../index.js";
+import type { Image, Tabular, Text, TokenizedText } from "#dataset/index";
+import type { NormalizedImage } from "#processing/image";
 
 /**
  * The data & label format goes through various stages.
@@ -37,7 +32,7 @@ type Token = number;
  * prediction needs data input and outputs label
  **/
 export interface ModelEncoded {
-  image: [image: processing.NormalizedImage<3>, label: number];
+  image: [image: NormalizedImage<3>, label: number];
   tabular: [row: List<number>, number];
   text: [line: TokenizedText, next: Token];
 }
diff --git a/discojs/src/types/datatype.ts b/discojs/src/types/datatype.ts
new file mode 100644
index 000000000..8c7f6dbc7
--- /dev/null
+++ b/discojs/src/types/datatype.ts
@@ -0,0 +1,2 @@
+export const dataTypeValues = ["image", "tabular", "text"] as const;
+export type DataType = (typeof dataTypeValues)[number];
diff --git a/discojs/src/types/index.ts b/discojs/src/types/index.ts
index a8f8ef7fa..d7c33c215 100644
--- a/discojs/src/types/index.ts
+++ b/discojs/src/types/index.ts
@@ -1,5 +1,7 @@
+// eslint-disable-next-line no-restricted-syntax -- namespace re-export acceptable here
 export * as DataFormat from "./data_format.js";
 
-export const dataTypeValues = ["image", "tabular", "text"] as const;
-export type DataType = (typeof dataTypeValues)[number];
-export type Network = "decentralized" | "federated" | "local";
+export { dataTypeValues } from "./datatype.js";
+export type { DataType } from "./datatype.js";
+
+export type { Network } from "./network.js";
diff --git a/discojs/src/types/network.ts b/discojs/src/types/network.ts
new file mode 100644
index 000000000..0a27da3c6
--- /dev/null
+++ b/discojs/src/types/network.ts
@@ -0,0 +1 @@
+export type Network = "decentralized" | "federated" | "local";
diff --git a/discojs/src/utils/async_iterator.spec.ts b/discojs/src/utils/async_iterator.spec.ts
index 5b7dae9a0..e8aec5e8a 100644
--- a/discojs/src/utils/async_iterator.spec.ts
+++ b/discojs/src/utils/async_iterator.spec.ts
@@ -1,5 +1,5 @@
 import { describe, expect, it } from "vitest";
-import { gather, split } from "./async_iterator.js";
+import { gather, split } from "#utils/async_iterator";
 
 // Array.fromAsync not yet widely used (2024)
 async function arrayFromAsync<T>(iter: AsyncIterable<T>): Promise<T[]> {
diff --git a/discojs/src/validator.ts b/discojs/src/validator.ts
index a1bd01c39..ac493e960 100644
--- a/discojs/src/validator.ts
+++ b/discojs/src/validator.ts
@@ -1,12 +1,12 @@
-import type {
-  Dataset,
-  DataFormat,
-  DataType,
-  Model,
-  Task,
-  Network,
-} from "./index.js";
-import { processing } from "./index.js";
+import type { DataFormat, DataType, Network } from "#types/index";
+import type { Task } from "#task/index";
+import type { Dataset } from "#dataset/index";
+import type { Model } from "#models/index";
+import {
+  preprocess,
+  preprocessWithoutLabel,
+  postprocess,
+} from "#processing/index";
 
 export class Validator<D extends DataType> {
   readonly #model: Model<D>;
@@ -22,7 +22,7 @@ export class Validator<D extends DataType> {
   test(
     dataset: Dataset<DataFormat.Raw[D]>,
   ): Dataset<Record<"predicted" | "truth", DataFormat.Inferred[D]>> {
-    const preprocessed = processing.preprocess(this.task, dataset);
+    const preprocessed = preprocess(this.task, dataset);
     const batched = preprocessed.batch(this.task.trainingInformation.batchSize);
 
     const predictionWithTruth = batched
@@ -34,8 +34,8 @@ export class Validator<D extends DataType> {
       .flatten();
 
     return predictionWithTruth.map(([predicted, truth]) => ({
-      predicted: processing.postprocess(this.task, predicted),
-      truth: processing.postprocess(this.task, truth),
+      predicted: postprocess(this.task, predicted),
+      truth: postprocess(this.task, truth),
     }));
   }
 
@@ -43,14 +43,13 @@ export class Validator<D extends DataType> {
   async *infer(
     dataset: Dataset<DataFormat.RawWithoutLabel[D]>,
   ): AsyncGenerator<DataFormat.Inferred[D], void> {
-    const modelPredictions = processing
-      .preprocessWithoutLabel(this.task, dataset)
+    const modelPredictions = preprocessWithoutLabel(this.task, dataset)
       .batch(this.task.trainingInformation.batchSize)
       .map((batch) => this.#model.predict(batch))
       .flatten();
 
     const predictions = modelPredictions.map((prediction) =>
-      processing.postprocess(this.task, prediction),
+      postprocess(this.task, prediction),
     );
 
     for await (const e of predictions) yield e;
diff --git a/discojs/src/weights/aggregation.spec.ts b/discojs/src/weights/aggregation.spec.ts
index b9757be71..78d2e39ab 100644
--- a/discojs/src/weights/aggregation.spec.ts
+++ b/discojs/src/weights/aggregation.spec.ts
@@ -1,9 +1,10 @@
 import { assert, describe, it } from "vitest";
-import { WeightsContainer, aggregation } from "./index.js";
+import { WeightsContainer } from "#weights/weights_container";
+import { avg, sum, diff } from "#weights/aggregation";
 
 describe("weights aggregation", () => {
   it("avg of weights with two operands", () => {
-    const actual = aggregation.avg([
+    const actual = avg([
       WeightsContainer.of([1, 2, 3, -1], [-5, 6]),
       WeightsContainer.of([2, 3, 7, 1], [-10, 5]),
       WeightsContainer.of([3, 1, 5, 3], [-15, 19]),
@@ -14,7 +15,7 @@ describe("weights aggregation", () => {
   });
 
   it("sum of weights with two operands", () => {
-    const actual = aggregation.sum([
+    const actual = sum([
       [[3, -4], [9]],
       [[2, 13], [0]],
     ]);
@@ -24,7 +25,7 @@ describe("weights aggregation", () => {
   });
 
   it("diff of weights with two operands", () => {
-    const actual = aggregation.diff([
+    const actual = diff([
       [
         [3, -4, 5],
         [9, 1],
diff --git a/discojs/src/weights/aggregation.ts b/discojs/src/weights/aggregation.ts
index 868d23b8a..58773d13e 100644
--- a/discojs/src/weights/aggregation.ts
+++ b/discojs/src/weights/aggregation.ts
@@ -1,8 +1,8 @@
 import { List } from "immutable";
 import * as tf from "@tensorflow/tfjs";
 
-import type { TensorLike } from "./weights_container.js";
-import { WeightsContainer } from "./weights_container.js";
+import type { TensorLike } from "#weights/weights_container";
+import { WeightsContainer } from "#weights/weights_container";
 
 type WeightsLike = Iterable<TensorLike>;
 
diff --git a/discojs/src/weights/index.ts b/discojs/src/weights/index.ts
index b44515510..15df42a0a 100644
--- a/discojs/src/weights/index.ts
+++ b/discojs/src/weights/index.ts
@@ -1,2 +1,2 @@
 export { WeightsContainer } from "./weights_container.js";
-export * as aggregation from "./aggregation.js";
+export { sum, avg } from "./aggregation.js";
diff --git a/discojs/tsconfig.json b/discojs/tsconfig.json
index 3a4772a71..7b7cf6dca 100644
--- a/discojs/tsconfig.json
+++ b/discojs/tsconfig.json
@@ -1,7 +1,4 @@
 {
-  "compilerOptions": {
-    "composite": true
-  },
   "files": [],
   "references": [
     {
diff --git a/discojs/tsconfig.lib.json b/discojs/tsconfig.lib.json
index 90ec2cefa..6d137af39 100644
--- a/discojs/tsconfig.lib.json
+++ b/discojs/tsconfig.lib.json
@@ -1,10 +1,24 @@
 {
   "extends": "../tsconfig.base.lib.json",
   "compilerOptions": {
+    // Copied from `package.json` since some tools don't work without it
+    "paths": {
+      "#root/*": ["./src/*"],
+      "#aggregator/*": ["./src/aggregator/*"],
+      "#client/*": ["./src/client/*"],
+      "#dataset/*": ["./src/dataset/*"],
+      "#types/*": ["./src/types/*"],
+      "#logging/*": ["./src/logging/*"],
+      "#models/*": ["./src/models/*"],
+      "#processing/*": ["./src/processing/*"],
+      "#serialization/*": ["./src/serialization/*"],
+      "#task/*": ["./src/task/*"],
+      "#training/*": ["./src/training/*"],
+      "#utils/*": ["./src/utils/*"],
+      "#weights/*": ["./src/weights/*"]
+    },
     "rootDir": "./src",
-    "outDir": "dist",
-    "composite": true
+    "outDir": "dist"
   },
-  "include": ["src"],
-  "exclude": ["**/*.spec.ts"]
+  "include": ["src"]
 }
diff --git a/discojs/tsconfig.vitest.json b/discojs/tsconfig.vitest.json
index 7060e216b..f33d5a873 100644
--- a/discojs/tsconfig.vitest.json
+++ b/discojs/tsconfig.vitest.json
@@ -1,8 +1,7 @@
 {
   "extends": "../tsconfig.base.json",
   "compilerOptions": {
-    "noEmit": true,
-    "composite": true
+    "noEmit": true
   },
   "include": ["src"]
 }
diff --git a/docs/examples/custom_task.ts b/docs/examples/custom_task.ts
index feb3f8e29..7e8074a8d 100644
--- a/docs/examples/custom_task.ts
+++ b/docs/examples/custom_task.ts
@@ -1,7 +1,7 @@
 import tf from "@tensorflow/tfjs-node";
 
 import type { TaskProvider, ModelCard } from "@epfml/discojs";
-import { defaultTasks, defaultModels, models } from "@epfml/discojs";
+import { defaultTasks, defaultModels, TFJS } from "@epfml/discojs";
 import { Server as DiscoServer } from "server";
 
 // Define your own model card
@@ -32,7 +32,7 @@ const customModelCard: ModelCard<"tabular"> = {
       metrics: ["accuracy"],
     });
 
-    return Promise.resolve(new models.TFJS("tabular", model));
+    return Promise.resolve(new TFJS("tabular", model));
   },
 };
 
diff --git a/docs/examples/tsconfig.json b/docs/examples/tsconfig.json
index e1d75b158..b421cc14b 100644
--- a/docs/examples/tsconfig.json
+++ b/docs/examples/tsconfig.json
@@ -12,13 +12,6 @@
     }
   ],
   "compilerOptions": {
-    "module": "node16",
-    "target": "es2022",
-
-    "strict": true,
-
-    "skipLibCheck": true,
-
     "outDir": "dist",
     "rootDir": ".",
     "types": ["node"]
diff --git a/docs/examples/wikitext.ts b/docs/examples/wikitext.ts
index 6bf3e06e8..742417da6 100644
--- a/docs/examples/wikitext.ts
+++ b/docs/examples/wikitext.ts
@@ -1,6 +1,7 @@
 import "@tensorflow/tfjs-node";
 
-import { Disco, fetchTasks, models, Task } from "@epfml/discojs";
+import type { GPT, Task } from "@epfml/discojs";
+import { Disco, fetchTasks } from "@epfml/discojs";
 import {
   saveModelToDisk,
   loadModelFromDisk,
@@ -37,14 +38,12 @@ async function main(): Promise<void> {
     await disco.trainFully(dataset);
 
     // Get the model and save the trained model
-    model = disco.trainer.model as models.GPT;
+    model = disco.trainer.model as GPT;
     await saveModelToDisk(model, modelFolder, modelFileName);
     await disco.close();
   } else {
     // Load the trained model
-    model = (await loadModelFromDisk(
-      `${modelFolder}/${modelFileName}`,
-    )) as models.GPT;
+    model = (await loadModelFromDisk(`${modelFolder}/${modelFileName}`)) as GPT;
   }
 
   // Preprocess prompt
diff --git a/eslint.config.js b/eslint.config.js
index 72b6e8cf2..e5c9f3119 100644
--- a/eslint.config.js
+++ b/eslint.config.js
@@ -1,5 +1,4 @@
 // @ts-check
-
 import pluginVitest from "@vitest/eslint-plugin";
 import skipFormatting from "@vue/eslint-config-prettier/skip-formatting";
 import {
@@ -50,6 +49,7 @@ export default defineConfigWithVueTs(
           typesToIgnore: ["Model", "DataType"],
         },
       ],
+      "@typescript-eslint/consistent-type-imports": "error",
     },
   },
   {
@@ -59,6 +59,55 @@ export default defineConfigWithVueTs(
       "@typescript-eslint/no-unnecessary-type-assertion": "off",
     },
   },
+  {
+    // allow relative imports only in index.ts barrel files in discojs
+    files: ["discojs/**"],
+    ignores: ["**/index.ts"],
+    rules: {
+      "no-restricted-imports": [
+        "error",
+        {
+          patterns: [
+            {
+              group: ["./*", "../*"],
+              message:
+                "Use alias imports ('#/...') instead of relative imports.",
+            },
+            {
+              group: ["#/index"],
+              message:
+                "Do not import the root barrel from within src. Import the specific subfolder barrel instead (#/subfolder/index).",
+            },
+          ],
+        },
+      ],
+    },
+  },
+  {
+    // barrel files in discojs should contain only re-exports, and no `export *`
+    files: ["discojs/**/index.ts"],
+    rules: {
+      "no-restricted-syntax": [
+        "error",
+        {
+          selector:
+            "Program > :not(ExportNamedDeclaration, ExportAllDeclaration)",
+          message: "Barrel files may only re-export.",
+        },
+        {
+          selector: "ExportNamedDeclaration:not([source])",
+          message: "Barrel files may only re-export.",
+        },
+        {
+          selector: "ExportAllDeclaration",
+          message: "No `export *`. Use explicit named re-exports.",
+        },
+      ],
+    },
+  },
+  {
+    files: ["src/**/*.ts"],
+  },
   {
     ...pluginVitest.configs.recommended,
     files: ["**/*.spec.ts"],
diff --git a/onnx-converter/package.json b/onnx-converter/package.json
index 17d6eeb1a..9bd9679d0 100644
--- a/onnx-converter/package.json
+++ b/onnx-converter/package.json
@@ -6,7 +6,7 @@
   "scripts": {
     "convert_onnx": "pnpm run build && node dist/convert_onnx.js",
     "watch": "nodemon --ext ts --ignore dist --exec pnpm run",
-    "build": "tsc --build tsconfig.lib.json && cp -r src/protobuf dist",
+    "build": "tsc --build && cp -r src/protobuf dist",
     "lint": "pnpm exec eslint .",
     "test": ": nothing"
   },
diff --git a/onnx-converter/src/convert_onnx.ts b/onnx-converter/src/convert_onnx.ts
index fb00d0edd..505860bac 100644
--- a/onnx-converter/src/convert_onnx.ts
+++ b/onnx-converter/src/convert_onnx.ts
@@ -3,7 +3,7 @@ import { Map, Range } from "immutable";
 import fsPromise from "node:fs/promises";
 import * as tf from "@tensorflow/tfjs-node";
 
-import { models, serialization } from "@epfml/discojs";
+import { GPT, modelEncode } from "@epfml/discojs";
 
 const OUTPUT_FILENAME = "model.json";
 const GPT2_N_LAYER = 12;
@@ -33,7 +33,7 @@ async function main() {
 
   // Init empty TF.js model
   // Context length value from https://huggingface.co/Xenova/gpt2/blob/main/config.json
-  const gptModel = new models.GPT({ modelType: "gpt2", contextLength: 1024 });
+  const gptModel = new GPT({ modelType: "gpt2", contextLength: 1024 });
   if (gptModel.config.nLayer != GPT2_N_LAYER)
     throw new Error(
       `ONNX conversion only supports GPT-2 with 12 layers, instead found ${gptModel.config.nLayer}.`,
@@ -82,7 +82,7 @@ async function main() {
 
   gptLayersModel.setWeights(finalWeights); // shape or transpose mismatch will throw here
 
-  const encoded = await serialization.model.encode(gptModel);
+  const encoded = await modelEncode(gptModel);
   await fsPromise.writeFile(OUTPUT_FILENAME, encoded);
   console.log(`GPT-TFJS model saved to ${OUTPUT_FILENAME}`);
 }
diff --git a/onnx-converter/tsconfig.json b/onnx-converter/tsconfig.json
index abb5cd6c2..ba8a8381b 100644
--- a/onnx-converter/tsconfig.json
+++ b/onnx-converter/tsconfig.json
@@ -1,7 +1,4 @@
 {
-  "compilerOptions": {
-    "composite": true
-  },
   "files": [],
   "references": [
     {
diff --git a/package.json b/package.json
index 2c0fdc52c..5d8c1e385 100644
--- a/package.json
+++ b/package.json
@@ -4,7 +4,8 @@
   "scripts": {
     "lint": "eslint",
     "format:check": "prettier -c .",
-    "format:fix": "prettier -w ."
+    "format:fix": "prettier -w .",
+    "check_cycles": "dpdm --tsconfig discojs/tsconfig.lib.json --circular --exit-code circular:1 'discojs/src/**/*.ts'"
   },
   "dependencies": {
     "@tensorflow/tfjs-node": "catalog:"
@@ -15,6 +16,7 @@
     "@vitest/eslint-plugin": "1.6.12",
     "@vue/eslint-config-prettier": "10.2.0",
     "@vue/eslint-config-typescript": "14.9.0",
+    "dpdm": "4.3.0",
     "eslint": "10.0.3",
     "eslint-plugin-cypress": "6.2.0",
     "eslint-plugin-vue": "10.9.2",
diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml
index 363d1a0b5..0937a0e65 100644
--- a/pnpm-lock.yaml
+++ b/pnpm-lock.yaml
@@ -250,6 +250,9 @@ importers:
       '@vue/eslint-config-typescript':
         specifier: 14.9.0
         version: 14.9.0(eslint-plugin-vue@10.9.2(@typescript-eslint/parser@8.62.1(eslint@10.0.3(jiti@2.7.0)(supports-color@8.1.1))(supports-color@8.1.1)(typescript@6.0.3))(eslint@10.0.3(jiti@2.7.0)(supports-color@8.1.1))(vue-eslint-parser@10.4.1(eslint@10.0.3(jiti@2.7.0)(supports-color@8.1.1))(supports-color@8.1.1)))(eslint@10.0.3(jiti@2.7.0)(supports-color@8.1.1))(supports-color@8.1.1)(typescript@6.0.3)
+      dpdm:
+        specifier: 4.3.0
+        version: 4.3.0
       eslint:
         specifier: 10.0.3
         version: 10.0.3(jiti@2.7.0)(supports-color@8.1.1)
@@ -2347,6 +2350,10 @@ packages:
     resolution: {integrity: sha512-oKnbhFyRIXpUuez8iBMmyEa4nbj4IOQyuhc/wy9kY7/WVPcwIO9VA668Pu8RkO7+0G76SLROeyw9CpQ061i4mA==}
     engines: {node: '>=10'}
 
+  chalk@5.6.2:
+    resolution: {integrity: sha512-7NzBL0rN6fMUW+f7A6Io4h40qQlG+xGmtMxfbnH/K7TAtt8JQWVQK+6g0UXKMeVJoyV5EkkNsErQ8pVD3bLHbA==}
+    engines: {node: ^12.17.0 || ^14.13 || >=16.0.0}
+
   chart.js@4.5.1:
     resolution: {integrity: sha512-GIjfiT9dbmHRiYi6Nl2yFCq7kkwdkp1W/lp2J99rX0yo9tgJGn3lKQATztIjb5tVtevcBtIdICNWqlq5+E8/Pw==}
     engines: {pnpm: '>=8'}
@@ -2378,6 +2385,10 @@ packages:
     resolution: {integrity: sha512-aCj4O5wKyszjMmDT4tZj93kxyydN/K5zPWSCe6/0AV/AA1pqe5ZBIw0a2ZfPQV7lL5/yb5HsUreJ6UFAF1tEQw==}
     engines: {node: '>=18'}
 
+  cli-spinners@3.4.0:
+    resolution: {integrity: sha512-bXfOC4QcT1tKXGorxL3wbJm6XJPDqEnij2gQ2m7ESQuE+/z9YFIWnl/5RpTiKWbMq3EVKR4fRLJGn6DVfu0mpw==}
+    engines: {node: '>=18.20'}
+
   cli-table3@0.6.1:
     resolution: {integrity: sha512-w0q/enDHhPLq44ovMGdQeeDLvwxwavsJX7oQGYt/LrBlYsyaxyDnp6z3QzFut/6kLLKnlcUVJLrpB7KBfgG/RA==}
     engines: {node: 10.* || >= 12.*}
@@ -2389,6 +2400,10 @@ packages:
   cliui@7.0.4:
     resolution: {integrity: sha512-OcRE68cOsVMXp1Yvonl/fzkQOyjLSu/8bhPDfQt0e0/Eb283TKP20Fs2MqoPsr9SwA595rRCA+QMzYc9nBP+JQ==}
 
+  cliui@9.0.1:
+    resolution: {integrity: sha512-k7ndgKhwoQveBL+/1tqGJYNz097I7WOvwbmmU2AR5+magtbjPWQTS1C5vzGkBC8Ym8UWRzfKUzUUqFLypY4Q+w==}
+    engines: {node: '>=20'}
+
   color-convert@1.9.3:
     resolution: {integrity: sha512-QfAUtd+vFdAtFQcC8CCyYt1fYWxSqAiK2cSD6zDB8N3cpsEBAvRxp9zOGg6G/SHHJYAT88/az/IuDGALsNVbGg==}
 
@@ -2719,6 +2734,10 @@ packages:
     engines: {node: '>=12'}
     deprecated: Use your platform's native DOMException instead
 
+  dpdm@4.3.0:
+    resolution: {integrity: sha512-2ZrP5B3MHHo7mXWgNxntU5DhJIvXNP4hcMBdCJsuxEWenKMdf4h6XouuKMf3Sj2SAUNBy/ajcWAVIyiOZly6rw==}
+    hasBin: true
+
   driver.js@1.4.0:
     resolution: {integrity: sha512-Gm64jm6PmcU+si21sQhBrTAM1JvUrR0QhNmjkprNLxohOBzul9+pNHXgQaT9lW84gwg9GMLB3NZGuGolsz5uew==}
 
@@ -3070,6 +3089,10 @@ packages:
   fs-constants@1.0.0:
     resolution: {integrity: sha512-y6OAwoSIf7FyjMIv94u+b5rdheZEjzR63GTyZJm5qh4Bi+2YgwLCcI/fPFZkL5PSixOt6ZNKm+w+Hfp/Bciwow==}
 
+  fs-extra@11.4.0:
+    resolution: {integrity: sha512-EQsFzMUJkCKGr1ePqlYADkIUmHW1s3ZXr5Yqy6wbGrfUCphpl2maM/kyOIRA2HpP3AaFQTZXD4ldjek+nccddA==}
+    engines: {node: '>=14.14'}
+
   fs-extra@9.1.0:
     resolution: {integrity: sha512-hcg3ZmepS30/7BSFqRvoo3DOMQu7IjqxO5nCDt+zM9XWjb33Wg7ziNT+Qvqbuc3+gWpzO02JubVyk2G4Zvo1OQ==}
     engines: {node: '>=10'}
@@ -3143,6 +3166,10 @@ packages:
     deprecated: Old versions of glob are not supported, and contain widely publicized security vulnerabilities, which have been fixed in the current version. Please update. Support for old versions may be purchased (at exorbitant rates) by contacting i@izs.me
     hasBin: true
 
+  glob@13.0.6:
+    resolution: {integrity: sha512-Wjlyrolmm8uDpm/ogGyXZXb1Z+Ca2B8NbJwqBVg0axK9GbBeoS7yGV6vjXnYdGm6X53iehEuxxbyiKp8QmN4Vw==}
+    engines: {node: 18 || 20 || >=22}
+
   glob@7.2.3:
     resolution: {integrity: sha512-nFR0zLpU2YCaRxwoCJvL6UvCH2JFyFVIvwTLsIf21AuHlMskA1hhTdk+LlYJtOlYt9v6dvszD2BGRqBL+iQK9Q==}
     deprecated: Old versions of glob are not supported, and contain widely publicized security vulnerabilities, which have been fixed in the current version. Please update. Support for old versions may be purchased (at exorbitant rates) by contacting i@izs.me
@@ -3304,6 +3331,10 @@ packages:
     resolution: {integrity: sha512-iwGqO3J21aaSkC7jWnHP/difazwS7SFeIqxv6wEtLU8Y5KlzFTjyqcSIT0d8s4+dDhKytsk9PJZ2BkS5eZwQRQ==}
     engines: {node: '>=10'}
 
+  is-interactive@2.0.0:
+    resolution: {integrity: sha512-qP1vozQRI+BMOPcjFzrjXuQvdak2pHNUMZoeG2eRbiSqyvbEf/wQtEOTOX1guk6E3t36RkaqiSt8A/6YElNxLQ==}
+    engines: {node: '>=12'}
+
   is-number@7.0.0:
     resolution: {integrity: sha512-41Cifkg6e8TylSpdtTpeLVMqvSBEVzTttHvERD741+pnZ8ANv0004MRL43QKPDlK9cGvNp6NZWZUBlbGXYxxng==}
     engines: {node: '>=0.12.0'}
@@ -3329,6 +3360,10 @@ packages:
     resolution: {integrity: sha512-knxG2q4UC3u8stRGyAVJCOdxFmv5DZiRcdlIaAQXAbSfJya+OhopNotLQrstBhququ4ZpuKbDc/8S6mgXgPFPw==}
     engines: {node: '>=10'}
 
+  is-unicode-supported@2.1.0:
+    resolution: {integrity: sha512-mE00Gnza5EEB3Ds0HfMyllZzbBrmLOX3vfWoj9A9PEnTfratQ/BcaJOuMhnkhjXvb2+FkY3VuHqtAGpTPmglFQ==}
+    engines: {node: '>=18'}
+
   is-what@5.5.0:
     resolution: {integrity: sha512-oG7cgbmg5kLYae2N5IVd3jm2s+vldjxJzK1pcu9LfpGuQ93MQSzo0okvRna+7y5ifrD+20FE8FvjusyGaz14fw==}
     engines: {node: '>=18'}
@@ -3528,6 +3563,10 @@ packages:
     resolution: {integrity: sha512-8XPvpAA8uyhfteu8pIvQxpJZ7SYYdpUivZpGy6sFsBuKRY/7rQGavedeB8aK+Zkyq6upMFVL/9AW6vOYzfRyLg==}
     engines: {node: '>=10'}
 
+  log-symbols@7.0.1:
+    resolution: {integrity: sha512-ja1E3yCr9i/0hmBVaM0bfwDjnGy8I/s6PP4DFp+yP+a+mrHO4Rm7DtmnqROTUkHIkqffC84YY7AeqX6oFk0WFg==}
+    engines: {node: '>=18'}
+
   log-update@6.1.0:
     resolution: {integrity: sha512-9ie8ItPR6tjY5uYJh8K/Zrv/RMZ5VOlOWvtZdEHYSTFKZfIBPQa9tOAEeAWhd+AnIneLJ22w5fjOYtoutpWq5w==}
     engines: {node: '>=18'}
@@ -3781,6 +3820,10 @@ packages:
     resolution: {integrity: sha512-6IpQ7mKUxRcZNLIObR0hz7lxsapSSIYNZJwXPGeF0mTVqGKFIXj1DQcMoT22S3ROcLyY/rz0PWaWZ9ayWmad9g==}
     engines: {node: '>= 0.8.0'}
 
+  ora@9.4.1:
+    resolution: {integrity: sha512-6VlU9MLXbjVQD04AZCMX28hVtA5bUoadvUqO76MUCVA0ilwJbMiHsITRPfyVm6p/BC0Av/BXMujx39WCe1LEqw==}
+    engines: {node: '>=20'}
+
   ospath@1.2.2:
     resolution: {integrity: sha512-o6E5qJV5zkAbIDNhGSIlyOhScKXgQrSRMilfph0clDfM0nEnBOlKlH4sWDmG95BW/CvwNz0vmm7dJVtU2KlMiA==}
 
@@ -3840,6 +3883,10 @@ packages:
     resolution: {integrity: sha512-Xa4Nw17FS9ApQFJ9umLiJS4orGjm7ZzwUrwamcGQuHSzDyth9boKDaycYdDcZDuqYATXw4HFXgaqWTctW/v1HA==}
     engines: {node: '>=16 || 14 >=14.18'}
 
+  path-scurry@2.0.2:
+    resolution: {integrity: sha512-3O/iVVsJAPsOnpwWIeD+d6z/7PmqApyQePUtCndjatj/9I5LylHvt5qluFaBT3I5h3r1ejfR056c+FCv+NnNXg==}
+    engines: {node: 18 || 20 || >=22}
+
   path-to-regexp@8.4.2:
     resolution: {integrity: sha512-qRcuIdP69NPm4qbACK+aDogI5CBDMi1jKe0ry5rSQJz8JVLsC7jV8XpiJjGRLLol3N+R5ihGYcrPLTno6pAdBA==}
 
@@ -4222,6 +4269,10 @@ packages:
   std-env@4.0.0:
     resolution: {integrity: sha512-zUMPtQ/HBY3/50VbpkupYHbRroTRZJPRLvreamgErJVys0ceuzMkD44J/QjqhHjOzK42GQ3QZIeFG1OYfOtKqQ==}
 
+  stdin-discarder@0.3.2:
+    resolution: {integrity: sha512-eCPu1qRxPVkl5605OTWF8Wz40b4Mf45NY5LQmVPQ599knfs5QhASUm9GbJ5BDMDOXgrnh0wyEdvzmL//YMlw0A==}
+    engines: {node: '>=18'}
+
   streamx@2.28.0:
     resolution: {integrity: sha512-1Yowhzjf0ivGMrTIkY9hav5TxobO9qIVqUE41fiCGMGgc3CLlf4MY+9AHmZqBWgDTue0fY9zWjYFVyf6Diuobw==}
 
@@ -4463,6 +4514,11 @@ packages:
       eslint: ^8.57.0 || ^9.0.0 || ^10.0.0
       typescript: '>=4.8.4 <6.1.0'
 
+  typescript@5.9.3:
+    resolution: {integrity: sha512-jl1vZzPDinLr9eUt3J/t7V6FgNEw9QjvBPdysz9KfQDD41fQrC2Y4vKQdiaUpFT4bXlb1RHhLpp8wtm6M5TgSw==}
+    engines: {node: '>=14.17'}
+    hasBin: true
+
   typescript@6.0.3:
     resolution: {integrity: sha512-y2TvuxSZPDyQakkFRPZHKFm+KKVqIisdg9/CZwm9ftvKXLP8NRWj38/ODjNbr43SsoXqNuAisEf1GdCxqWcdBw==}
     engines: {node: '>=14.17'}
@@ -4818,10 +4874,18 @@ packages:
     resolution: {integrity: sha512-y11nGElTIV+CT3Zv9t7VKl+Q3hTQoT9a1Qzezhhl6Rp21gJ/IVTW7Z3y9EWXhuUBC2Shnf+DX0antecpAwSP8w==}
     engines: {node: '>=10'}
 
+  yargs-parser@22.0.0:
+    resolution: {integrity: sha512-rwu/ClNdSMpkSrUb+d6BRsSkLUq1fmfsY6TOpYzTwvwkg1/NRG85KBy3kq++A8LKQwX6lsu+aWad+2khvuXrqw==}
+    engines: {node: ^20.19.0 || ^22.12.0 || >=23}
+
   yargs@16.2.2:
     resolution: {integrity: sha512-Nt9ZJjXTv5R8MHbqby/wXQ6Gi0Bb3TcYZkR1bzuL4yB2OxWPkXknz513gEF0GoA6tn00UpbPvERW8rzCuWCA6w==}
     engines: {node: '>=10'}
 
+  yargs@18.1.0:
+    resolution: {integrity: sha512-2rAgRKu54VsHkqI0/tYkmluGXHD4KW7yZoycuqDQ15QOTnc2VVfy0nN/1eMhnQLO00A+dwtK20xuCnc1YGeUyg==}
+    engines: {node: ^20.19.0 || ^22.12.0 || >=23}
+
   yauzl@3.4.0:
     resolution: {integrity: sha512-jIH9yLR9wqr0wOS0TpBvo/g/2UgZH5qePVbjgRliiF0BYvOZyaBknKsF+x9Iht0O6sqgnB93rCICdOZFecJuDw==}
     engines: {node: '>=12'}
@@ -4830,6 +4894,10 @@ packages:
     resolution: {integrity: sha512-rVksvsnNCdJ/ohGc6xgPwyN8eheCxsiLM8mxuE/t/mOVqJewPuO1miLpTHQiRgTKCLexL4MeAFVagts7HmNZ2Q==}
     engines: {node: '>=10'}
 
+  yoctocolors@2.2.0:
+    resolution: {integrity: sha512-xYqdZFUK/VYazNl/oCDYN+3WloWQwMfZxBoiNt6qNyk+xfOdi598muWE42rNZFp1kNOiqW936q5RhUdnpqElSg==}
+    engines: {node: '>=18'}
+
   zod@3.25.76:
     resolution: {integrity: sha512-gzUt/qt81nXsFGKIFcC3YnfEAx5NkunCfnDlvuBSSFS02bcXu4Lmea0AFIUwbLWxWPx3d9p8S5QoaujKcNQxcQ==}
 
@@ -6501,6 +6569,8 @@ snapshots:
       ansi-styles: 4.3.0
       supports-color: 7.2.0
 
+  chalk@5.6.2: {}
+
   chart.js@4.5.1:
     dependencies:
       '@kurkle/color': 0.3.4
@@ -6533,6 +6603,8 @@ snapshots:
     dependencies:
       restore-cursor: 5.1.0
 
+  cli-spinners@3.4.0: {}
+
   cli-table3@0.6.1:
     dependencies:
       string-width: 4.2.3
@@ -6550,6 +6622,12 @@ snapshots:
       strip-ansi: 6.0.1
       wrap-ansi: 7.0.0
 
+  cliui@9.0.1:
+    dependencies:
+      string-width: 7.2.0
+      strip-ansi: 7.2.0
+      wrap-ansi: 9.0.2
+
   color-convert@1.9.3:
     dependencies:
       color-name: 1.1.3
@@ -6914,6 +6992,16 @@ snapshots:
       webidl-conversions: 7.0.0
     optional: true
 
+  dpdm@4.3.0:
+    dependencies:
+      chalk: 5.6.2
+      fs-extra: 11.4.0
+      glob: 13.0.6
+      ora: 9.4.1
+      tslib: 2.8.1
+      typescript: 5.9.3
+      yargs: 18.1.0
+
   driver.js@1.4.0: {}
 
   dunder-proto@1.0.1:
@@ -7301,6 +7389,12 @@ snapshots:
 
   fs-constants@1.0.0: {}
 
+  fs-extra@11.4.0:
+    dependencies:
+      graceful-fs: 4.2.11
+      jsonfile: 6.2.1
+      universalify: 2.0.1
+
   fs-extra@9.1.0:
     dependencies:
       at-least-node: 1.0.0
@@ -7388,6 +7482,12 @@ snapshots:
       package-json-from-dist: 1.0.1
       path-scurry: 1.11.1
 
+  glob@13.0.6:
+    dependencies:
+      minimatch: 10.2.5
+      minipass: 7.1.3
+      path-scurry: 2.0.2
+
   glob@7.2.3:
     dependencies:
       fs.realpath: 1.0.0
@@ -7530,6 +7630,8 @@ snapshots:
       global-dirs: 3.0.1
       is-path-inside: 3.0.3
 
+  is-interactive@2.0.0: {}
+
   is-number@7.0.0: {}
 
   is-path-inside@3.0.3: {}
@@ -7544,6 +7646,8 @@ snapshots:
 
   is-unicode-supported@0.1.0: {}
 
+  is-unicode-supported@2.1.0: {}
+
   is-what@5.5.0: {}
 
   isexe@2.0.0: {}
@@ -7749,6 +7853,11 @@ snapshots:
       chalk: 4.1.2
       is-unicode-supported: 0.1.0
 
+  log-symbols@7.0.1:
+    dependencies:
+      is-unicode-supported: 2.1.0
+      yoctocolors: 2.2.0
+
   log-update@6.1.0:
     dependencies:
       ansi-escapes: 7.3.0
@@ -7971,6 +8080,17 @@ snapshots:
       type-check: 0.4.0
       word-wrap: 1.2.5
 
+  ora@9.4.1:
+    dependencies:
+      chalk: 5.6.2
+      cli-cursor: 5.0.0
+      cli-spinners: 3.4.0
+      is-interactive: 2.0.0
+      is-unicode-supported: 2.1.0
+      log-symbols: 7.0.1
+      stdin-discarder: 0.3.2
+      string-width: 8.2.1
+
   ospath@1.2.2: {}
 
   outdent@0.8.0: {}
@@ -8059,6 +8179,11 @@ snapshots:
       lru-cache: 10.4.3
       minipass: 7.1.3
 
+  path-scurry@2.0.2:
+    dependencies:
+      lru-cache: 11.5.1
+      minipass: 7.1.3
+
   path-to-regexp@8.4.2: {}
 
   pathe@2.0.3: {}
@@ -8526,6 +8651,8 @@ snapshots:
 
   std-env@4.0.0: {}
 
+  stdin-discarder@0.3.2: {}
+
   streamx@2.28.0:
     dependencies:
       events-universal: 1.0.1
@@ -8792,6 +8919,8 @@ snapshots:
     transitivePeerDependencies:
       - supports-color
 
+  typescript@5.9.3: {}
+
   typescript@6.0.3: {}
 
   typical@4.0.0: {}
@@ -9074,6 +9203,8 @@ snapshots:
 
   yargs-parser@20.2.9: {}
 
+  yargs-parser@22.0.0: {}
+
   yargs@16.2.2:
     dependencies:
       cliui: 7.0.4
@@ -9084,12 +9215,23 @@ snapshots:
       y18n: 5.0.8
       yargs-parser: 20.2.9
 
+  yargs@18.1.0:
+    dependencies:
+      cliui: 9.0.1
+      escalade: 3.2.0
+      get-caller-file: 2.0.5
+      string-width: 8.2.1
+      y18n: 5.0.8
+      yargs-parser: 22.0.0
+
   yauzl@3.4.0:
     dependencies:
       pend: 1.2.0
 
   yocto-queue@0.1.0: {}
 
+  yoctocolors@2.2.0: {}
+
   zod@3.25.76: {}
 
   zod@4.3.6: {}
diff --git a/server/package.json b/server/package.json
index 4a3f29bdf..0d6c60b70 100644
--- a/server/package.json
+++ b/server/package.json
@@ -11,7 +11,7 @@
   "scripts": {
     "watch": "nodemon --ext ts --ignore dist --watch ../discojs-node/dist --watch . --exec pnpm run",
     "start": "node dist/main.js",
-    "build": "tsc --build tsconfig.lib.json",
+    "build": "tsc --build",
     "test": "cd .. && vitest --run --project=server"
   },
   "author": "",
diff --git a/server/src/controllers/decentralized_controller.ts b/server/src/controllers/decentralized_controller.ts
index 33bdc403a..f35672392 100644
--- a/server/src/controllers/decentralized_controller.ts
+++ b/server/src/controllers/decentralized_controller.ts
@@ -4,12 +4,13 @@ import * as msgpack from "@msgpack/msgpack";
 import type WebSocket from "ws";
 import { Map } from "immutable";
 
-import { client, DataType } from "@epfml/discojs";
+import type { DataType, NodeID } from "@epfml/discojs";
+import { mtype, decentralizedMessages } from "@epfml/discojs";
 
 import { TrainingController } from "./training_controller.js";
 
-import messages = client.decentralized.messages;
-import MessageTypes = client.messages.type;
+import messages = decentralizedMessages;
+import MessageTypes = mtype.MType;
 
 const debug = createDebug("server:controllers:decentralized");
 
@@ -20,7 +21,7 @@ export class DecentralizedController<
   // The boolean value indicates if the node is ready to exchange weight updates (i.e.
   // the node has already sent a PeerIsReady message)
   // We wait for all peers to be ready to exchange weight updates
-  #roundPeers = Map<client.NodeID, boolean>();
+  #roundPeers = Map<NodeID, boolean>();
   #aggregationRound = 0;
 
   handle(ws: WebSocket): void {
@@ -143,7 +144,7 @@ export class DecentralizedController<
         debug("Sending peer list to: %o", id.slice(0, 4));
 
         const encoded = msgpack.encode(readyPeerIDs);
-        return [id, encoded] as [client.NodeID, Buffer];
+        return [id, encoded] as [NodeID, Buffer];
       })
       .map(([id, encoded]) => {
         const conn = this.connections.get(id);
diff --git a/server/src/controllers/federated_controller.ts b/server/src/controllers/federated_controller.ts
index 68d017bea..d4d0fdfd2 100644
--- a/server/src/controllers/federated_controller.ts
+++ b/server/src/controllers/federated_controller.ts
@@ -1,19 +1,20 @@
 import createDebug from "debug";
-import WebSocket from "ws";
+import type WebSocket from "ws";
 import { v4 as randomUUID } from "uuid";
 import * as msgpack from "@msgpack/msgpack";
 
-import type { DataType, Task } from "@epfml/discojs";
+import type { DataType, Task, Encoded } from "@epfml/discojs";
 import {
-  aggregator as aggregators,
-  client,
-  serialization,
+  mtype,
+  federatedMessages,
+  weightsEncode,
+  weightsDecode,
+  MeanAggregator,
 } from "@epfml/discojs";
 
 import { TrainingController } from "./training_controller.js";
 
-import MessageTypes = client.messages.type;
-import FederatedMessages = client.federated.messages;
+import MessageTypes = mtype.MType;
 
 const debug = createDebug("server:controllers:federated");
 
@@ -25,25 +26,24 @@ export class FederatedController<D extends DataType> extends TrainingController<
    * Aggregators for each hosted task.
     By default the server waits for 100% of the nodes to send their contributions before aggregating the updates
    */
-  #aggregator = new aggregators.MeanAggregator(undefined, 1, "relative");
+  #aggregator = new MeanAggregator(undefined, 1, "relative");
   /**
    * The most up to date global weights. The model weights are already serialized and
    * can be sent to participants, before starting training, or when joining mid-training
    * or staled participants
    */
-  #latestGlobalWeights: serialization.Encoded;
+  #latestGlobalWeights: Encoded;
 
   constructor(
     task: Task<D, "federated">,
-    private readonly initialWeights: serialization.Encoded,
+    private readonly initialWeights: Encoded,
   ) {
     super(task);
     this.#latestGlobalWeights = this.initialWeights;
 
     // Save the latest weight updates to be able to send it to new or outdated clients
     this.#aggregator.on("aggregation", async (weightUpdate) => {
-      this.#latestGlobalWeights =
-        await serialization.weights.encode(weightUpdate);
+      this.#latestGlobalWeights = await weightsEncode(weightUpdate);
     });
   }
 
@@ -71,7 +71,7 @@ export class FederatedController<D extends DataType> extends TrainingController<
     // Setup callbacks triggered upon receiving the different client messages
     ws.on("message", (data: Buffer) => {
       const msg: unknown = msgpack.decode(data);
-      if (!FederatedMessages.isMessageFederated(msg)) {
+      if (!federatedMessages.isMessageFederated(msg)) {
         debug("invalid federated message received on WebSocket: %o", msg);
         return; // TODO send back error
       }
@@ -87,7 +87,7 @@ export class FederatedController<D extends DataType> extends TrainingController<
           debug(`client [%s] joined ${this.task.id}`, shortId);
           this.connections = this.connections.set(clientId, ws); // add the new client
 
-          const msg: FederatedMessages.NewFederatedNodeInfo = {
+          const msg: federatedMessages.NewFederatedNodeInfo = {
             type: MessageTypes.NewFederatedNodeInfo,
             id: clientId,
             waitForMoreParticipants:
@@ -107,7 +107,7 @@ export class FederatedController<D extends DataType> extends TrainingController<
         case MessageTypes.SendPayload: {
           const { payload, round } = msg;
           if (this.#aggregator.isValidContribution(clientId, round)) {
-            const weights = serialization.weights.decode(payload);
+            const weights = weightsDecode(payload);
 
             // Create a callback to send the aggregated weight to the client
             // when enough contributions are received
@@ -117,10 +117,10 @@ export class FederatedController<D extends DataType> extends TrainingController<
                 this.#aggregator.round,
                 shortId,
               );
-              const msg: FederatedMessages.ReceiveServerPayload = {
+              const msg: federatedMessages.ReceiveServerPayload = {
                 type: MessageTypes.ReceiveServerPayload,
                 round: this.#aggregator.round, // send the current round number after aggregation
-                payload: await serialization.weights.encode(weightUpdate),
+                payload: await weightsEncode(weightUpdate),
                 nbOfParticipants: this.connections.size,
               };
               ws.send(msgpack.encode(msg));
@@ -142,7 +142,7 @@ export class FederatedController<D extends DataType> extends TrainingController<
             // no latest model at the first round
             if (this.#latestGlobalWeights === undefined) return;
 
-            const msg: FederatedMessages.ReceiveServerPayload = {
+            const msg: federatedMessages.ReceiveServerPayload = {
               type: MessageTypes.ReceiveServerPayload,
               round: this.#aggregator.round - 1, // send the model from the previous round
               payload: this.#latestGlobalWeights,
@@ -165,11 +165,7 @@ export class FederatedController<D extends DataType> extends TrainingController<
       // Reset the training session when all participants left
       if (this.connections.size === 0) {
         debug("All participants left. Resetting the training session");
-        this.#aggregator = new aggregators.MeanAggregator(
-          undefined,
-          1,
-          "relative",
-        );
+        this.#aggregator = new MeanAggregator(undefined, 1, "relative");
         this.#latestGlobalWeights = this.initialWeights;
       }
 
diff --git a/server/src/controllers/training_controller.ts b/server/src/controllers/training_controller.ts
index ff3e52c6f..6db7d2fb1 100644
--- a/server/src/controllers/training_controller.ts
+++ b/server/src/controllers/training_controller.ts
@@ -3,8 +3,8 @@ import type WebSocket from "ws";
 import { Map } from "immutable";
 import * as msgpack from "@msgpack/msgpack";
 
-import { client } from "@epfml/discojs";
-import type { DataType, Network, Task } from "@epfml/discojs";
+import { mtype } from "@epfml/discojs";
+import type { DataType, Network, Task, NodeID } from "@epfml/discojs";
 
 const debug = createDebug("server:controllers");
 
@@ -37,7 +37,7 @@ export abstract class TrainingController<
    * the list allows updating participants about the training status
    * i.e. waiting for more participants or resuming training
    */
-  protected connections = Map<client.NodeID, WebSocket>();
+  protected connections = Map<NodeID, WebSocket>();
 
   constructor(protected readonly task: Task<D, N>) {}
 
@@ -48,7 +48,7 @@ export abstract class TrainingController<
    *
    * @param currentId the id of the participant that just joined
    */
-  protected sendEnoughParticipantsMsgIfNeeded(currentId: client.NodeID) {
+  protected sendEnoughParticipantsMsgIfNeeded(currentId: NodeID) {
     // If we are currently waiting for more participants to join and we now have enough,
     // broadcast to previously waiting participants that the training can start
     if (
@@ -64,8 +64,8 @@ export abstract class TrainingController<
             "Sending enough-participant message to client [%s]",
             participantId.slice(0, 4),
           );
-          const msg: client.messages.EnoughParticipants = {
-            type: client.messages.type.EnoughParticipants,
+          const msg: mtype.EnoughParticipants = {
+            type: mtype.MType.EnoughParticipants,
             nbOfParticipants: this.connections.size,
           };
           participantWs.send(msgpack.encode(msg));
@@ -86,8 +86,8 @@ export abstract class TrainingController<
         "Telling remaining client [%s] to wait for participants",
         participantId.slice(0, 4),
       );
-      const msg: client.messages.WaitingForMoreParticipants = {
-        type: client.messages.type.WaitingForMoreParticipants,
+      const msg: mtype.WaitingForMoreParticipants = {
+        type: mtype.MType.WaitingForMoreParticipants,
         nbOfParticipants: this.connections.size,
       };
       participantWs.send(msgpack.encode(msg));
diff --git a/server/src/model_set.ts b/server/src/model_set.ts
index 021f1b7c2..ad2fc9aca 100644
--- a/server/src/model_set.ts
+++ b/server/src/model_set.ts
@@ -1,10 +1,16 @@
 import { Map } from "immutable";
 import "@tensorflow/tfjs-node";
 
-import type { DataType, ModelCardInfo, ModelCard } from "@epfml/discojs";
-import { EventEmitter, Model, serialization } from "@epfml/discojs";
+import type {
+  DataType,
+  ModelCardInfo,
+  ModelCard,
+  Model,
+  Encoded,
+} from "@epfml/discojs";
+import { EventEmitter, modelEncode, isEncoded } from "@epfml/discojs";
 
-type EncodedModel = serialization.Encoded;
+type EncodedModel = Encoded;
 type AvailableModel = [ModelCardInfo<DataType>, EncodedModel];
 
 /**
@@ -77,13 +83,13 @@ export class ModelSet extends EventEmitter<{
     let encodedModel: EncodedModel;
     if (!Array.isArray(newModel)) {
       const model = await newModel.getModel();
-      encodedModel = await serialization.model.encode(model);
+      encodedModel = await modelEncode(model);
     } else {
       const model = newModel[1];
-      if (serialization.isEncoded(model)) {
+      if (isEncoded(model)) {
         encodedModel = model; // don't do anything if already encoded
       } else {
-        encodedModel = await serialization.model.encode(model);
+        encodedModel = await modelEncode(model);
       }
     }
 
diff --git a/server/src/routes/task_router.ts b/server/src/routes/task_router.ts
index 13e2dab31..b5b80bf8a 100644
--- a/server/src/routes/task_router.ts
+++ b/server/src/routes/task_router.ts
@@ -3,9 +3,13 @@ import type { Request, Response } from "express";
 import express from "express";
 import { Set } from "immutable";
 
-import type { DataType, ModelCardInfo, Network } from "@epfml/discojs";
+import type { DataType, ModelCardInfo, Network, Encoded } from "@epfml/discojs";
 import { Task } from "@epfml/discojs";
-import { serialization } from "@epfml/discojs";
+import {
+  serializeTaskToJSON,
+  deserializeTaskFromJSON,
+  modelDecode,
+} from "@epfml/discojs";
 
 import type { TaskSet } from "../task_set.js";
 import type { ModelSet } from "../model_set.js";
@@ -45,7 +49,7 @@ export class TaskRouter {
       res.status(200).send(
         this.#taskSet.tasks
           .valueSeq()
-          .map(([t, _]) => serialization.task.serializeToJSON(t))
+          .map(([t, _]) => serializeTaskToJSON(t))
           .toArray(),
       );
     });
@@ -70,7 +74,7 @@ export class TaskRouter {
               z.string(),
               z.array(z.number()).transform((bytes) => Uint8Array.from(bytes)),
             ]),
-            task: z.any().transform(serialization.task.deserializeFromJSON),
+            task: z.any().transform(deserializeTaskFromJSON),
           })
           .safeParseAsync(req.body);
 
@@ -129,11 +133,11 @@ export class TaskRouter {
    */
   private async registerUploadedModel(
     task: Task<DataType, Network>,
-    encoded: serialization.Encoded,
+    encoded: Encoded,
   ): Promise<ModelCardInfo.ID> {
     let uploaded;
     try {
-      uploaded = await serialization.model.decode(encoded);
+      uploaded = await modelDecode(encoded);
     } catch (e) {
       debug("posted model isn't a valid encoded model: %o", e);
       throw new Error("uploaded model is invalid");
diff --git a/server/src/routes/training_router.ts b/server/src/routes/training_router.ts
index 28da1e932..1c75387ba 100644
--- a/server/src/routes/training_router.ts
+++ b/server/src/routes/training_router.ts
@@ -1,11 +1,11 @@
 import express from "express";
 import type expressWS from "express-ws";
-import type { Task, DataType, Network } from "@epfml/discojs";
-import { serialization } from "@epfml/discojs";
+import type { Task, DataType, Network, Encoded } from "@epfml/discojs";
+import { modelDecode, weightsEncode } from "@epfml/discojs";
 
 import type { TaskSet } from "../task_set.js";
+import type { TrainingController } from "../controllers/index.js";
 import {
-  TrainingController,
   FederatedController,
   DecentralizedController,
 } from "../controllers/index.js";
@@ -44,7 +44,7 @@ export class TrainingRouter<N extends Exclude<Network, "local">> {
   // websocket connections
   private async onNewTask<D extends DataType>(
     task: Task<D, N>,
-    encodedModel: serialization.Encoded,
+    encodedModel: Encoded,
   ): Promise<void> {
     // The controller handles the actual logic of collaborative training
     // in its `handle` method. Each task has a dedicated controller which
@@ -55,10 +55,8 @@ export class TrainingRouter<N extends Exclude<Network, "local">> {
 
       // The federated controller takes the initial model weights at initialization
       // so that it can send it to new clients
-      const model = serialization.model.decode(encodedModel);
-      const encodedWeights = await serialization.weights.encode(
-        (await model).weights,
-      );
+      const model = modelDecode(encodedModel);
+      const encodedWeights = await weightsEncode((await model).weights);
       taskController = new FederatedController(t, encodedWeights);
     } else {
       const t = task as Task<D, "decentralized">;
diff --git a/server/src/task_set.ts b/server/src/task_set.ts
index 28ccfbaa4..10faedfb2 100644
--- a/server/src/task_set.ts
+++ b/server/src/task_set.ts
@@ -1,11 +1,11 @@
 import { Map } from "immutable";
 import "@tensorflow/tfjs-node";
 
-import type { DataType, Network, Task } from "@epfml/discojs";
-import { EventEmitter, serialization } from "@epfml/discojs";
-import { ModelSet } from "./model_set.js";
+import type { DataType, Network, Task, Encoded } from "@epfml/discojs";
+import { EventEmitter } from "@epfml/discojs";
+import type { ModelSet } from "./model_set.js";
 
-type EncodedModel = serialization.Encoded;
+type EncodedModel = Encoded;
 type TaskAndModel = [Task<DataType, Network>, EncodedModel];
 
 /**
diff --git a/server/tests/client.spec.ts b/server/tests/client.spec.ts
index ccdfd40e3..4568aad8e 100644
--- a/server/tests/client.spec.ts
+++ b/server/tests/client.spec.ts
@@ -6,8 +6,9 @@ import type {
   ModelCard,
 } from "@epfml/discojs";
 import {
-  aggregator as aggregators,
-  client as clients,
+  MeanAggregator,
+  DecentralizedClient,
+  FederatedClient,
   defaultTasks,
   defaultModels,
 } from "@epfml/discojs";
@@ -42,10 +43,10 @@ describe("decentralized client", () => {
       [defaultTasks.cifar10],
     );
 
-    const client = new clients.decentralized.DecentralizedClient(
+    const client = new DecentralizedClient(
       url,
       await defaultTasks.cifar10.getTask(),
-      new aggregators.MeanAggregator(),
+      new MeanAggregator(),
     );
 
     await client.connect();
@@ -55,10 +56,10 @@ describe("decentralized client", () => {
   it("fails to connect to invalid task", async () => {
     const url = await startServer([], []); // no models or tasks
 
-    const client = new clients.decentralized.DecentralizedClient(
+    const client = new DecentralizedClient(
       url,
       await defaultTasks.cifar10.getTask(),
-      new aggregators.MeanAggregator(),
+      new MeanAggregator(),
     );
 
     await expect(client.connect()).rejects.toThrow();
@@ -93,10 +94,10 @@ describe("federated client", () => {
       [defaultTasks.titanic],
     );
 
-    const client = new clients.federated.FederatedClient(
+    const client = new FederatedClient(
       url,
       await defaultTasks.titanic.getTask(),
-      new aggregators.MeanAggregator(),
+      new MeanAggregator(),
     );
 
     await client.connect();
@@ -106,10 +107,10 @@ describe("federated client", () => {
   it("fails to connect to invalid task", async () => {
     const url = await startServer([], []); // no task
 
-    const client = new clients.federated.FederatedClient(
+    const client = new FederatedClient(
       url,
       await defaultTasks.titanic.getTask(),
-      new aggregators.MeanAggregator(),
+      new MeanAggregator(),
     );
 
     await expect(client.connect()).rejects.toThrow();
diff --git a/server/tests/e2e/decentralized.spec.ts b/server/tests/e2e/decentralized.spec.ts
index 87059efdc..15e6a181e 100644
--- a/server/tests/e2e/decentralized.spec.ts
+++ b/server/tests/e2e/decentralized.spec.ts
@@ -2,13 +2,15 @@ import type * as http from "node:http";
 import type {
   DataType,
   RoundStatus,
+  Client,
   Task,
   TaskProvider,
   ModelCard,
 } from "@epfml/discojs";
 import {
-  aggregator as aggregators,
-  client as clients,
+  MeanAggregator,
+  SecureAggregator,
+  DecentralizedClient,
   Disco,
   defaultTasks,
   defaultModels,
@@ -69,18 +71,14 @@ describe("end-to-end decentralized", { timeout: 50_000 }, () => {
     aggregatorType: "mean" | "secure",
     input: number[],
     rounds: number,
-  ): Promise<[WeightsContainer, clients.Client<"decentralized">]> {
+  ): Promise<[WeightsContainer, Client<"decentralized">]> {
     const task = await defaultTasks.cifar10.getTask();
     const aggregator =
       aggregatorType === "mean"
-        ? new aggregators.MeanAggregator(0, 1, "relative")
-        : new aggregators.SecureAggregator();
+        ? new MeanAggregator(0, 1, "relative")
+        : new SecureAggregator();
 
-    const client = new clients.decentralized.DecentralizedClient(
-      url,
-      task,
-      aggregator,
-    );
+    const client = new DecentralizedClient(url, task, aggregator);
     await client.connect();
 
     // Perform multiple training rounds
diff --git a/server/tests/e2e/federated.spec.ts b/server/tests/e2e/federated.spec.ts
index 042d83cee..f6d3e2f5e 100644
--- a/server/tests/e2e/federated.spec.ts
+++ b/server/tests/e2e/federated.spec.ts
@@ -8,8 +8,9 @@ import type {
   Task,
   TaskProvider,
   WeightsContainer,
+  ModelCard,
 } from "@epfml/discojs";
-import { Disco, defaultTasks, defaultModels, ModelCard } from "@epfml/discojs";
+import { Disco, defaultTasks, defaultModels } from "@epfml/discojs";
 import { List } from "immutable";
 import { assert, afterEach, describe, expect, it } from "vitest";
 import { Server } from "../../src/index.js";
diff --git a/server/tests/routes.spec.ts b/server/tests/routes.spec.ts
index e8c0222a5..e825a08c9 100644
--- a/server/tests/routes.spec.ts
+++ b/server/tests/routes.spec.ts
@@ -15,9 +15,11 @@ import {
   defaultTasks,
   fetchModels,
   fetchTasks,
-  models,
   pushTask,
-  serialization,
+  modelDecode,
+  modelEncode,
+  serializeTaskToJSON,
+  TFJS,
 } from "@epfml/discojs";
 
 import { Server } from "../src/index.js";
@@ -228,8 +230,8 @@ describe("GET /tasks/:id/model.json", { timeout: 20_000 }, () => {
     expect(res.status).toBe(200);
     const encoded = new Uint8Array(await res.arrayBuffer());
     expect(encoded.length).toBeGreaterThan(0);
-    const model = await serialization.model.decode(encoded);
-    expect(model).toBeInstanceOf(models.TFJS);
+    const model = await modelDecode(encoded);
+    expect(model).toBeInstanceOf(TFJS);
   });
 
   it("answers 404 for an unknown task", async () => {
@@ -275,7 +277,7 @@ describe("POST /tasks", { timeout: 20_000 }, () => {
     const task = await newTask();
 
     const res = await postTask(url, {
-      task: serialization.task.serializeToJSON(task),
+      task: serializeTaskToJSON(task),
       model: defaultModels.TitanicClassifier.card.id,
     });
 
@@ -291,9 +293,7 @@ describe("POST /tasks", { timeout: 20_000 }, () => {
     );
 
     const res = await postTask(url, {
-      task: serialization.task.serializeToJSON(
-        await defaultTasks.titanic.getTask(),
-      ),
+      task: serializeTaskToJSON(await defaultTasks.titanic.getTask()),
       model: defaultModels.TitanicClassifier.card.id,
     });
 
@@ -306,12 +306,12 @@ describe("POST /tasks", { timeout: 20_000 }, () => {
       [defaultTasks.titanic],
     );
     const task = await newTask();
-    const encoded = await serialization.model.encode(
+    const encoded = await modelEncode(
       await defaultModels.TitanicClassifier.getModel(),
     );
 
     const res = await postTask(url, {
-      task: serialization.task.serializeToJSON(task),
+      task: serializeTaskToJSON(task),
       model: [...encoded],
     });
 
@@ -335,7 +335,7 @@ describe("POST /tasks", { timeout: 20_000 }, () => {
       [defaultModels.TitanicClassifier],
       [defaultTasks.titanic],
     );
-    const task = serialization.task.serializeToJSON(await newTask());
+    const task = serializeTaskToJSON(await newTask());
 
     // no model at all
     expect((await postTask(url, { task })).status).toBe(400);
@@ -350,12 +350,12 @@ describe("POST /tasks", { timeout: 20_000 }, () => {
     );
     // an image model, while the task is tabular. The mismatch is caught from
     // the model itself, nothing in the request states its data type
-    const encoded = await serialization.model.encode(
+    const encoded = await modelEncode(
       await defaultModels.LUSClassifier.getModel(),
     );
 
     const res = await postTask(url, {
-      task: serialization.task.serializeToJSON(await newTask()),
+      task: serializeTaskToJSON(await newTask()),
       model: [...encoded],
     });
 
@@ -373,7 +373,7 @@ describe("POST /tasks", { timeout: 20_000 }, () => {
     );
 
     const res = await postTask(url, {
-      task: serialization.task.serializeToJSON(await newTask()),
+      task: serializeTaskToJSON(await newTask()),
       model: "not-a-model",
     });
 
@@ -387,7 +387,7 @@ describe("POST /tasks", { timeout: 20_000 }, () => {
     );
 
     const res = await postTask(url, {
-      task: serialization.task.serializeToJSON(await newTask()), // tabular
+      task: serializeTaskToJSON(await newTask()), // tabular
       model: defaultModels.LUSClassifier.card.id, // image
     });
 
@@ -416,7 +416,7 @@ describe("POST /tasks", { timeout: 20_000 }, () => {
     const task = await newTask();
 
     const posted = await postTask(url, {
-      task: serialization.task.serializeToJSON(task),
+      task: serializeTaskToJSON(task),
       model: defaultModels.TitanicClassifier.card.id,
     });
     expect(posted.status).toBe(200);
diff --git a/server/tsconfig.json b/server/tsconfig.json
index 3a4772a71..7b7cf6dca 100644
--- a/server/tsconfig.json
+++ b/server/tsconfig.json
@@ -1,7 +1,4 @@
 {
-  "compilerOptions": {
-    "composite": true
-  },
   "files": [],
   "references": [
     {
diff --git a/server/tsconfig.lib.json b/server/tsconfig.lib.json
index 06455f74b..d1f865211 100644
--- a/server/tsconfig.lib.json
+++ b/server/tsconfig.lib.json
@@ -1,10 +1,5 @@
 {
   "extends": "../tsconfig.base.lib.json",
-  "compilerOptions": {
-    "rootDir": "./src",
-    "outDir": "dist",
-    "composite": true
-  },
   "references": [
     {
       "path": "../discojs/tsconfig.lib.json"
@@ -13,5 +8,9 @@
       "path": "../discojs-node/tsconfig.lib.json"
     }
   ],
+  "compilerOptions": {
+    "rootDir": "./src",
+    "outDir": "dist"
+  },
   "include": ["src"]
 }
diff --git a/server/tsconfig.vitest.json b/server/tsconfig.vitest.json
index 135cacb84..188d9e578 100644
--- a/server/tsconfig.vitest.json
+++ b/server/tsconfig.vitest.json
@@ -1,5 +1,15 @@
 {
   "extends": "../tsconfig.base.json",
-  "compilerOptions": { "noEmit": true },
+  "references": [
+    {
+      "path": "../discojs/tsconfig.lib.json"
+    },
+    {
+      "path": "../discojs-node/tsconfig.lib.json"
+    }
+  ],
+  "compilerOptions": {
+    "noEmit": true
+  },
   "include": ["src", "tests"]
 }
diff --git a/tsconfig.base.json b/tsconfig.base.json
index 8b5ab3610..0923fbcad 100644
--- a/tsconfig.base.json
+++ b/tsconfig.base.json
@@ -1,6 +1,6 @@
 {
-  "composite": true,
   "compilerOptions": {
+    "composite": true,
     "declaration": true,
     "declarationMap": true,
 
@@ -17,6 +17,11 @@
     "noEmitOnError": true,
 
     // enforce using the override keyword
-    "noImplicitOverride": true
+    "noImplicitOverride": true,
+
+    "isolatedModules": true,
+    "verbatimModuleSyntax": true,
+
+    "customConditions": ["@disco/source"]
   }
 }
diff --git a/tsconfig.base.lib.json b/tsconfig.base.lib.json
index e2c16c267..fc21f755e 100644
--- a/tsconfig.base.lib.json
+++ b/tsconfig.base.lib.json
@@ -1,6 +1,4 @@
 {
   "extends": "./tsconfig.base.json",
-  "compilerOptions": {
-    "customConditions": ["@disco/source"]
-  }
+  "exclude": ["**/*.spec.ts"]
 }
diff --git a/tsconfig.eslint.json b/tsconfig.eslint.json
deleted file mode 100644
index 1dc1c369a..000000000
--- a/tsconfig.eslint.json
+++ /dev/null
@@ -1,4 +0,0 @@
-{
-  "extends": "./tsconfig.base.lib.json",
-  "include": ["**/src", "**/tests", "**/*config.ts"]
-}
diff --git a/vitest.config.ts b/vitest.config.ts
index 536c237da..50aceec44 100644
--- a/vitest.config.ts
+++ b/vitest.config.ts
@@ -1,9 +1,16 @@
 import { defineConfig } from "vitest/config";
 
 export default defineConfig({
+  resolve: {
+    conditions: ["@disco/source"],
+  },
+  ssr: {
+    resolve: {
+      conditions: ["@disco/source"],
+    },
+  },
   test: {
     setupFiles: "./testSetupImportTFJSNode.ts",
-
     projects: [
       {
         extends: true,
diff --git a/webapp/cypress/e2e/task-creation.cy.ts b/webapp/cypress/e2e/task-creation.cy.ts
index 309404415..dbad26f4c 100644
--- a/webapp/cypress/e2e/task-creation.cy.ts
+++ b/webapp/cypress/e2e/task-creation.cy.ts
@@ -1,4 +1,4 @@
-import { serialization, type Task } from "@epfml/discojs";
+import { deserializeTaskFromJSON, type Task } from "@epfml/discojs";
 
 import * as tf from "@tensorflow/tfjs";
 
@@ -64,7 +64,7 @@ it("submits with tabular task", () => {
 
   cy.wait("@posted")
     .its("request.body.task")
-    .then(serialization.task.deserializeFromJSON)
+    .then(deserializeTaskFromJSON)
     .should("deep.equal", {
       id: "id",
       dataType: "tabular",
diff --git a/webapp/cypress/support/e2e.ts b/webapp/cypress/support/e2e.ts
index 119b229bc..0a55ba067 100644
--- a/webapp/cypress/support/e2e.ts
+++ b/webapp/cypress/support/e2e.ts
@@ -5,8 +5,9 @@ import type {
   Task,
   TaskProvider,
   TrainingInformation,
+  Encoded,
 } from "@epfml/discojs";
-import { serialization } from "@epfml/discojs";
+import { serializeTaskToJSON, modelEncode } from "@epfml/discojs";
 
 export function setupServerWith(
   ...providers: (Task<DataType, Network> | TaskProvider<DataType, Network>)[]
@@ -23,9 +24,7 @@ export function setupServerWith(
     .as("taskAndModels");
 
   cy.get<Array<[Task<DataType, Network>, unknown]>>("@taskAndModels")
-    .then((taskAndModels) =>
-      taskAndModels.map(([t]) => serialization.task.serializeToJSON(t)),
-    )
+    .then((taskAndModels) => taskAndModels.map(([t]) => serializeTaskToJSON(t)))
     .then((tasks) =>
       cy.intercept({ hostname: "server", pathname: "tasks" }, tasks),
     );
@@ -42,9 +41,7 @@ export function setupServerWith(
         { hostname: "server", pathname: `/tasks/${task.id}/model.json` },
         { statusCode: 200 },
       );
-      cy.wrap<Promise<serialization.Encoded>, serialization.Encoded>(
-        serialization.model.encode(model),
-      ).then((encoded) =>
+      cy.wrap<Promise<Encoded>, Encoded>(modelEncode(model)).then((encoded) =>
         cy.intercept(
           { hostname: "server", pathname: `/tasks/${task.id}/model.json` },
           (req) =>
diff --git a/webapp/cypress/tsconfig.json b/webapp/cypress/tsconfig.json
index 1640a71b8..639f355e8 100644
--- a/webapp/cypress/tsconfig.json
+++ b/webapp/cypress/tsconfig.json
@@ -1,5 +1,8 @@
 {
-  "extends": "@vue/tsconfig/tsconfig.dom.json",
+  "extends": [
+    "@vue/tsconfig/tsconfig.dom.json",
+    "../../tsconfig.base.lib.json"
+  ],
   "include": ["./e2e/**/*", "./support/**/*"],
   "references": [
     {
@@ -10,6 +13,9 @@
     }
   ],
   "compilerOptions": {
+    "module": "ESNext",
+    "moduleResolution": "Bundler",
+
     "isolatedModules": false,
     "target": "es2022",
     "lib": ["es2022", "dom"],
diff --git a/webapp/package.json b/webapp/package.json
index e54a9e5e6..00017ef77 100644
--- a/webapp/package.json
+++ b/webapp/package.json
@@ -4,7 +4,6 @@
   "type": "module",
   "scripts": {
     "start": "vite",
-    "prebuild": "pnpm -F webapp^... run build",
     "build": "vue-tsc --build && vite build",
     "test": "pnpm run test:unit && pnpm run test:e2e",
     "test:unit": "vitest --run",
diff --git a/webapp/src/components/containers/ButtonsCard.vue b/webapp/src/components/containers/ButtonsCard.vue
index 65666e7df..aeb0552cc 100644
--- a/webapp/src/components/containers/ButtonsCard.vue
+++ b/webapp/src/components/containers/ButtonsCard.vue
@@ -36,7 +36,7 @@
 <script setup lang="ts">
 import { computed } from "vue";
 
-import { List } from "immutable";
+import type { List } from "immutable";
 
 import CustomButton from "@/components/simple/CustomButton.vue";
 
diff --git a/webapp/src/components/containers/TableLayout.vue b/webapp/src/components/containers/TableLayout.vue
index 3548c0c93..a575480a8 100644
--- a/webapp/src/components/containers/TableLayout.vue
+++ b/webapp/src/components/containers/TableLayout.vue
@@ -23,7 +23,7 @@
 </template>
 
 <script lang="ts" setup>
-import { List } from "immutable";
+import type { List } from "immutable";
 
 defineProps<{
   columns: List<string>;
diff --git a/webapp/src/components/dataset_input/FileSelection.vue b/webapp/src/components/dataset_input/FileSelection.vue
index 12d75e0f6..4fae6dfd2 100644
--- a/webapp/src/components/dataset_input/FileSelection.vue
+++ b/webapp/src/components/dataset_input/FileSelection.vue
@@ -94,7 +94,8 @@
 </template>
 
 <script lang="ts" setup>
-import { Range, Set } from "immutable";
+import type { Set } from "immutable";
+import { Range } from "immutable";
 import { computed, ref } from "vue";
 
 import CustomButton from "@/components/simple/CustomButton.vue";
diff --git a/webapp/src/components/dataset_input/ImageDatasetInput.vue b/webapp/src/components/dataset_input/ImageDatasetInput.vue
index df9dd6a75..48e5b1426 100644
--- a/webapp/src/components/dataset_input/ImageDatasetInput.vue
+++ b/webapp/src/components/dataset_input/ImageDatasetInput.vue
@@ -5,7 +5,7 @@
 </template>
 
 <script setup lang="ts">
-import { Set } from "immutable";
+import type { Set } from "immutable";
 import { ref, watch } from "vue";
 
 import type { Image } from "@epfml/discojs";
diff --git a/webapp/src/components/dataset_input/LabeledImageDatasetInput/ByCSV.vue b/webapp/src/components/dataset_input/LabeledImageDatasetInput/ByCSV.vue
index 4b5831f2a..aaf21d62d 100644
--- a/webapp/src/components/dataset_input/LabeledImageDatasetInput/ByCSV.vue
+++ b/webapp/src/components/dataset_input/LabeledImageDatasetInput/ByCSV.vue
@@ -51,7 +51,8 @@
 
 <script lang="ts" setup>
 import * as d3 from "d3";
-import { Map, Set } from "immutable";
+import type { Set } from "immutable";
+import { Map } from "immutable";
 import { computed, ref, watch } from "vue";
 
 import { Dataset } from "@epfml/discojs";
diff --git a/webapp/src/components/dataset_input/LabeledImageDatasetInput/ByGroup.vue b/webapp/src/components/dataset_input/LabeledImageDatasetInput/ByGroup.vue
index f592b9385..ba53fd95b 100644
--- a/webapp/src/components/dataset_input/LabeledImageDatasetInput/ByGroup.vue
+++ b/webapp/src/components/dataset_input/LabeledImageDatasetInput/ByGroup.vue
@@ -11,7 +11,8 @@
 </template>
 
 <script lang="ts" setup>
-import { Map, Set } from "immutable";
+import type { Set } from "immutable";
+import { Map } from "immutable";
 import type { Ref, WatchHandle } from "vue";
 import { computed, ref, watch } from "vue";
 
diff --git a/webapp/src/components/dataset_input/LabeledImageDatasetInput/index.vue b/webapp/src/components/dataset_input/LabeledImageDatasetInput/index.vue
index bf5e2ed11..57c3857fe 100644
--- a/webapp/src/components/dataset_input/LabeledImageDatasetInput/index.vue
+++ b/webapp/src/components/dataset_input/LabeledImageDatasetInput/index.vue
@@ -47,7 +47,7 @@
 <script lang="ts" setup>
 defineOptions({ name: "LabeledImageDatasetInput" });
 
-import { Set } from "immutable";
+import type { Set } from "immutable";
 import { ref } from "vue";
 
 import DatasetInput from "../DatasetInput.vue";
diff --git a/webapp/src/components/dataset_input/TabularDatasetInput.vue b/webapp/src/components/dataset_input/TabularDatasetInput.vue
index 379ff527d..a7c7c7272 100644
--- a/webapp/src/components/dataset_input/TabularDatasetInput.vue
+++ b/webapp/src/components/dataset_input/TabularDatasetInput.vue
@@ -7,7 +7,7 @@
 </template>
 
 <script lang="ts" setup>
-import { Set } from "immutable";
+import type { Set } from "immutable";
 import { ref, watch } from "vue";
 
 import type { Dataset, Tabular } from "@epfml/discojs";
diff --git a/webapp/src/components/dataset_input/TextDatasetInput.vue b/webapp/src/components/dataset_input/TextDatasetInput.vue
index 8638b3b64..42350e1ba 100644
--- a/webapp/src/components/dataset_input/TextDatasetInput.vue
+++ b/webapp/src/components/dataset_input/TextDatasetInput.vue
@@ -7,11 +7,11 @@
 </template>
 
 <script lang="ts" setup>
-import { Set } from "immutable";
+import type { Set } from "immutable";
 import { ref, watch } from "vue";
 
 import type { Text } from "@epfml/discojs";
-import { Dataset } from "@epfml/discojs";
+import type { Dataset } from "@epfml/discojs";
 import { loadText } from "@epfml/discojs-web";
 
 import DatasetInput from "./DatasetInput.vue";
diff --git a/webapp/src/components/task_creation_form/TaskCreationForm.vue b/webapp/src/components/task_creation_form/TaskCreationForm.vue
index b45a8b026..fc6b0ce9d 100644
--- a/webapp/src/components/task_creation_form/TaskCreationForm.vue
+++ b/webapp/src/components/task_creation_form/TaskCreationForm.vue
@@ -616,7 +616,7 @@ import * as z from "zod";
 import * as tf from "@tensorflow/tfjs";
 
 import {
-  models,
+  TFJS,
   pushTask,
   Task,
   Tokenizer,
@@ -941,7 +941,7 @@ async function onSubmit(form: unknown): Promise<void> {
           loss,
           optimizer: tf.train[optimizer.name](optimizer.learningRate),
         });
-        model = new models.TFJS(task.dataType, loaded);
+        model = new TFJS(task.dataType, loaded);
         break;
       }
       case "text":
diff --git a/webapp/src/components/testing/ModelLibrary.vue b/webapp/src/components/testing/ModelLibrary.vue
index 88cfca446..e58fe9ec4 100644
--- a/webapp/src/components/testing/ModelLibrary.vue
+++ b/webapp/src/components/testing/ModelLibrary.vue
@@ -154,7 +154,7 @@ import { RouterLink } from "vue-router";
 import { VueSpinner } from "vue3-spinners";
 
 import type { DataType, Model, Network, Task } from "@epfml/discojs";
-import { client as clients, aggregator } from "@epfml/discojs";
+import { LocalClient, getAggregator } from "@epfml/discojs";
 
 import BinIcon from "@/assets/svg/BinIcon.vue";
 import { useToaster } from "@/composables/toaster";
@@ -240,7 +240,7 @@ async function downloadModel(task: Task<DataType, Network>): Promise<void> {
   try {
     toaster.info("Downloading model...");
 
-    const client = new clients.LocalClient(
+    const client = new LocalClient(
       CONFIG.serverUrl,
       {
         ...task,
@@ -249,7 +249,7 @@ async function downloadModel(task: Task<DataType, Network>): Promise<void> {
           scheme: "local",
         },
       } as Task<DataType, "local">,
-      aggregator.getAggregator(task),
+      getAggregator(task),
     );
     const model = await client.getLatestModel();
 
diff --git a/webapp/src/components/testing/__tests__/ModelLibrary.spec.ts b/webapp/src/components/testing/__tests__/ModelLibrary.spec.ts
index 5e92a8f63..d7b138e3e 100644
--- a/webapp/src/components/testing/__tests__/ModelLibrary.spec.ts
+++ b/webapp/src/components/testing/__tests__/ModelLibrary.spec.ts
@@ -7,12 +7,8 @@ import { flushPromises, mount } from "@vue/test-utils";
 import { createTestingPinia } from "@pinia/testing";
 import { createPersistedStatePlugin } from "pinia-plugin-persistedstate-2";
 
-import {
-  models as discoModels,
-  serialization,
-  Task,
-  Tokenizer,
-} from "@epfml/discojs";
+import type { Task } from "@epfml/discojs";
+import { GPT, serializeTaskToJSON, Tokenizer } from "@epfml/discojs";
 
 import { CONFIG } from "@/config";
 import { useModelsStore } from "@/store";
@@ -60,7 +56,7 @@ it("shows stored models", async () => {
   tasks.value = Map([[TASK.id, TASK]]);
 
   const models = useModelsStore();
-  await models.add("task", new discoModels.GPT());
+  await models.add("task", new GPT());
   await nextTick();
 
   expect(wrapper.get("div.text-xl").text()).to.equal("task title");
@@ -69,9 +65,7 @@ it("shows stored models", async () => {
 it("allows to download server's models", async () => {
   vi.stubGlobal("fetch", (url: string | URL) => {
     if (url.toString() === new URL("tasks", CONFIG.serverUrl).href)
-      return new Response(
-        JSON.stringify([serialization.task.serializeToJSON(TASK)]),
-      );
+      return new Response(JSON.stringify([serializeTaskToJSON(TASK)]));
     throw new Error(`unhandled get: ${url}`);
   });
   afterEach(() => {
diff --git a/webapp/src/components/training/TrainerDashboard.vue b/webapp/src/components/training/TrainerDashboard.vue
index dfd1d6c49..9324cdf00 100644
--- a/webapp/src/components/training/TrainerDashboard.vue
+++ b/webapp/src/components/training/TrainerDashboard.vue
@@ -151,7 +151,7 @@ import type {
   Task,
   Network,
 } from "@epfml/discojs";
-import { async_iterator, Disco } from "@epfml/discojs";
+import { split, Disco } from "@epfml/discojs";
 
 import { useToaster } from "@/composables/toaster";
 import TrainingInformation from "@/components/training/TrainingInformation.vue";
@@ -249,11 +249,11 @@ async function startTraining(): Promise<void> {
 
     roundsLogs.value = List<RoundLogs>();
     for await (const round of trainingGenerator.value) {
-      const [roundGen, roundLogs] = async_iterator.split(round);
+      const [roundGen, roundLogs] = split(round);
 
       roundGenerator.value = roundGen;
       for await (const epoch of roundGenerator.value) {
-        const [epochGen, epochLogs] = async_iterator.split(epoch);
+        const [epochGen, epochLogs] = split(epoch);
 
         epochGenerator.value = epochGen;
         for await (const batch of epochGenerator.value)
diff --git a/webapp/src/components/training/TrainingInformation.vue b/webapp/src/components/training/TrainingInformation.vue
index b1602c3e8..114eabcd3 100644
--- a/webapp/src/components/training/TrainingInformation.vue
+++ b/webapp/src/components/training/TrainingInformation.vue
@@ -161,7 +161,8 @@ import {
   type ChartOptions,
   type TooltipPositionerFunction,
 } from "chart.js";
-import { List, Range } from "immutable";
+import type { List } from "immutable";
+import { Range } from "immutable";
 import { computed, ref } from "vue";
 import { Line } from "vue-chartjs";
 
diff --git a/webapp/src/components/training/__tests__/TrainerDashboard.spec.ts b/webapp/src/components/training/__tests__/TrainerDashboard.spec.ts
index 2374369e2..a60d6aa59 100644
--- a/webapp/src/components/training/__tests__/TrainerDashboard.spec.ts
+++ b/webapp/src/components/training/__tests__/TrainerDashboard.spec.ts
@@ -1,6 +1,6 @@
 import * as fs from "node:fs/promises";
 import { CONFIG } from "@/config";
-import { defaultTasks, serialization } from "@epfml/discojs";
+import { defaultTasks, modelEncode } from "@epfml/discojs";
 import { loadCSV } from "@epfml/discojs-web";
 import { createTestingPinia } from "@pinia/testing";
 import { mount } from "@vue/test-utils";
@@ -18,7 +18,7 @@ async function setupForTask() {
       new URL("tasks/titanic/model.json", CONFIG.serverUrl).href
     ) {
       const model = await provider.modelCard.getModel();
-      const encoded = await serialization.model.encode(model);
+      const encoded = await modelEncode(model);
       return new Response(new Uint8Array(encoded));
     }
     throw new Error(`unhandled get: ${url}`);
diff --git a/webapp/src/store/models/__tests__/models.spec.ts b/webapp/src/store/models/__tests__/models.spec.ts
index 175e3b2f7..b390ee06e 100644
--- a/webapp/src/store/models/__tests__/models.spec.ts
+++ b/webapp/src/store/models/__tests__/models.spec.ts
@@ -3,7 +3,7 @@ import { createPersistedStatePlugin } from "pinia-plugin-persistedstate-2";
 import { beforeEach, expect, it } from "vitest";
 import { createApp } from "vue";
 
-import { models as discoModels } from "@epfml/discojs";
+import { GPT } from "@epfml/discojs";
 
 import { useModelsStore } from "../index";
 
@@ -17,7 +17,7 @@ beforeEach(() => {
 it("stores model", async () => {
   const models = useModelsStore();
 
-  const id = await models.add("task id", new discoModels.GPT());
+  const id = await models.add("task id", new GPT());
   expect(models.infos.size).to.equal(1);
-  expect(await models.get(id)).to.be.an.instanceof(discoModels.GPT);
+  expect(await models.get(id)).to.be.an.instanceof(GPT);
 });
diff --git a/webapp/src/store/models/index.ts b/webapp/src/store/models/index.ts
index aca613343..8ac462707 100644
--- a/webapp/src/store/models/index.ts
+++ b/webapp/src/store/models/index.ts
@@ -3,7 +3,7 @@ import { defineStore } from "pinia";
 import { computed, shallowRef, toRaw } from "vue";
 
 import type { DataType, Model } from "@epfml/discojs";
-import { serialization } from "@epfml/discojs";
+import { modelEncode, modelDecode } from "@epfml/discojs";
 
 import { bestStorage } from "./storage";
 import type { ModelID, State } from "./types";
@@ -29,7 +29,7 @@ export const useModelsStore = defineStore(
       const infos = idToModel.value.get(id);
       if (infos === undefined) return undefined;
 
-      return await serialization.model.decode(toRaw(infos.encoded));
+      return await modelDecode(toRaw(infos.encoded));
     }
 
     async function add(
@@ -42,7 +42,7 @@ export const useModelsStore = defineStore(
       idToModel.value = idToModel.value.set(id, {
         taskID,
         dateSaved,
-        encoded: await serialization.model.encode(model),
+        encoded: await modelEncode(model),
       });
 
       return id;
diff --git a/webapp/src/store/models/opfs.ts b/webapp/src/store/models/opfs.ts
index 32fbdf80f..51d50ef8d 100644
--- a/webapp/src/store/models/opfs.ts
+++ b/webapp/src/store/models/opfs.ts
@@ -2,7 +2,7 @@ import { Map } from "immutable";
 import * as msgpack from "@msgpack/msgpack";
 import type { IStorage } from "pinia-plugin-persistedstate-2";
 
-import { serialization } from "@epfml/discojs";
+import { isEncoded } from "@epfml/discojs";
 
 import type { Storage } from "./storage";
 import { UNSUPPORTED_STORAGE } from "./storage";
@@ -102,7 +102,7 @@ function isSerializedInfos(raw: unknown): raw is OPFS.Serialized[0][1] {
 
   if (
     typeof taskID !== "string" ||
-    !serialization.isEncoded(encoded) ||
+    !isEncoded(encoded) ||
     !(dateSaved instanceof Date)
   )
     return false;
diff --git a/webapp/src/store/models/types.ts b/webapp/src/store/models/types.ts
index 1b15ca242..aa63c7214 100644
--- a/webapp/src/store/models/types.ts
+++ b/webapp/src/store/models/types.ts
@@ -1,13 +1,13 @@
 import type { Map } from "immutable";
 
-import type { serialization } from "@epfml/discojs";
+import type { Encoded } from "@epfml/discojs";
 
 export type ModelID = number;
 
 export interface Infos {
   taskID: string;
   dateSaved: Date;
-  encoded: serialization.Encoded;
+  encoded: Encoded;
 }
 
 // only `ref` is `idToModel`
diff --git a/webapp/src/store/tasks.ts b/webapp/src/store/tasks.ts
index 4a27c5dec..b0d9806dd 100644
--- a/webapp/src/store/tasks.ts
+++ b/webapp/src/store/tasks.ts
@@ -1,5 +1,6 @@
 import createDebug from "debug";
-import { Map, Set } from "immutable";
+import type { Map } from "immutable";
+import { Set } from "immutable";
 import { defineStore } from "pinia";
 import { shallowRef } from "vue";
 
diff --git a/webapp/tsconfig.app.json b/webapp/tsconfig.app.json
index 6b13f7305..f654118ed 100644
--- a/webapp/tsconfig.app.json
+++ b/webapp/tsconfig.app.json
@@ -1,5 +1,5 @@
 {
-  "extends": "@vue/tsconfig/tsconfig.dom.json",
+  "extends": ["@vue/tsconfig/tsconfig.dom.json", "../tsconfig.base.lib.json"],
   "include": ["env.d.ts", "src/**/*", "src/**/*.vue"],
   "exclude": ["src/**/__tests__/*"],
   "references": [
@@ -11,7 +11,11 @@
     }
   ],
   "compilerOptions": {
+    "module": "ESNext",
+    "moduleResolution": "Bundler",
+
     "composite": true,
+    "customConditions": ["@disco/source"],
     "tsBuildInfoFile": "tsconfig.app.tsbuildinfo",
     "paths": {
       "@/*": ["./src/*"]
diff --git a/webapp/tsconfig.json b/webapp/tsconfig.json
index 42627fa65..6fdcd090b 100644
--- a/webapp/tsconfig.json
+++ b/webapp/tsconfig.json
@@ -13,8 +13,5 @@
     {
       "path": "./cypress/tsconfig.json"
     }
-  ],
-  "compilerOptions": {
-    "module": "NodeNext"
-  }
+  ]
 }
diff --git a/webapp/tsconfig.node.json b/webapp/tsconfig.node.json
index dfa46af11..f391327ff 100644
--- a/webapp/tsconfig.node.json
+++ b/webapp/tsconfig.node.json
@@ -1,5 +1,5 @@
 {
-  "extends": "@tsconfig/node20/tsconfig.json",
+  "extends": ["@tsconfig/node20/tsconfig.json", "../tsconfig.base.lib.json"],
   "include": [
     "vite.config.*",
     "vitest.config.*",
diff --git a/webapp/tsconfig.vitest.json b/webapp/tsconfig.vitest.json
index f2c24d817..c9279b454 100644
--- a/webapp/tsconfig.vitest.json
+++ b/webapp/tsconfig.vitest.json
@@ -10,7 +10,6 @@
     }
   ],
   "compilerOptions": {
-    "composite": true,
     "tsBuildInfoFile": "tsconfig.vitest.tsbuildinfo",
     "lib": [],
     "types": ["node", "jsdom"]
diff --git a/webapp/vite.config.ts b/webapp/vite.config.ts
index 58e28e24d..43f9787f6 100644
--- a/webapp/vite.config.ts
+++ b/webapp/vite.config.ts
@@ -1,6 +1,7 @@
 import { fileURLToPath, URL } from "node:url";
 import tailwindcss from "@tailwindcss/vite";
 import vue from "@vitejs/plugin-vue";
+import { defaultClientConditions, defaultServerConditions } from "vite";
 import { defineConfig } from "vite";
 
 // https://vitejs.dev/config/
@@ -8,6 +9,7 @@ export default defineConfig({
   server: { port: 1351 },
   plugins: [tailwindcss(), vue()],
   resolve: {
+    conditions: ["@disco/source", ...defaultClientConditions],
     alias: {
       "@": fileURLToPath(new URL("./src", import.meta.url)),
       // TODO until simple-peer#883 is fixed
@@ -15,6 +17,9 @@ export default defineConfig({
       buffer: "buffer/",
     },
   },
+  ssr: {
+    resolve: { conditions: ["@disco/source", ...defaultServerConditions] },
+  },
   optimizeDeps: {
     include: ["buffer"],
   },