From 001e7a86ff478b6e49aaee827687134eab1bc3ec Mon Sep 17 00:00:00 2001
From: tharvik <tharvik@users.noreply.github.com>
Date: Wed, 28 Feb 2024 13:23:13 +0100
Subject: [PATCH 1/2] discojs-core: don't reexport tfjs

---
 cli/src/data.ts                                   |  3 ++-
 discojs/discojs-core/src/aggregator/base.ts       |  5 +++--
 discojs/discojs-core/src/aggregator/mean.spec.ts  |  6 ++++--
 discojs/discojs-core/src/aggregator/mean.ts       |  6 ++++--
 discojs/discojs-core/src/aggregator/secure.ts     |  9 +++++----
 discojs/discojs-core/src/client/base.ts           | 12 +++++++-----
 discojs/discojs-core/src/dataset/data/data.ts     |  9 +++++----
 .../src/dataset/data/image_data.spec.ts           |  3 ++-
 .../discojs-core/src/dataset/data/image_data.ts   |  6 ++++--
 .../src/dataset/data/preprocessing/base.ts        | 10 ++++++----
 .../data/preprocessing/image_preprocessing.ts     |  7 ++++---
 .../data/preprocessing/tabular_preprocessing.ts   |  6 ++++--
 .../data/preprocessing/text_preprocessing.ts      |  7 ++++---
 .../src/dataset/data/tabular_data.spec.ts         |  3 ++-
 .../src/dataset/data_loader/image_loader.ts       | 10 ++++++----
 discojs/discojs-core/src/dataset/dataset.ts       |  2 +-
 discojs/discojs-core/src/default_tasks/cifar10.ts |  5 ++++-
 discojs/discojs-core/src/default_tasks/geotags.ts |  5 ++++-
 .../discojs-core/src/default_tasks/lus_covid.ts   |  5 ++++-
 discojs/discojs-core/src/default_tasks/mnist.ts   |  4 +++-
 .../discojs-core/src/default_tasks/simple_face.ts |  5 ++++-
 .../discojs-core/src/default_tasks/skin_mnist.ts  |  4 +++-
 discojs/discojs-core/src/default_tasks/titanic.ts |  5 ++++-
 discojs/discojs-core/src/index.ts                 |  2 --
 .../discojs-core/src/logging/trainer_logger.ts    |  2 +-
 discojs/discojs-core/src/memory/base.ts           |  6 +++---
 discojs/discojs-core/src/memory/empty.ts          |  5 +++--
 discojs/discojs-core/src/privacy.ts               |  4 +++-
 .../discojs-core/src/serialization/model.spec.ts  |  3 ++-
 discojs/discojs-core/src/serialization/model.ts   |  2 +-
 discojs/discojs-core/src/serialization/weights.ts |  3 ++-
 discojs/discojs-core/src/task/task_handler.ts     |  6 ++++--
 discojs/discojs-core/src/task/task_provider.ts    |  4 +++-
 .../src/training/trainer/distributed_trainer.ts   |  7 +++++--
 .../src/training/trainer/local_trainer.ts         |  3 ++-
 .../discojs-core/src/training/trainer/trainer.ts  |  8 ++++++--
 .../src/training/trainer/trainer_builder.ts       |  9 ++++++---
 .../src/training/training_functions.ts            |  4 +++-
 discojs/discojs-core/src/types.ts                 |  9 ++++++---
 discojs/discojs-core/src/validation/validator.ts  |  4 +++-
 discojs/discojs-core/src/weights/aggregation.ts   |  5 +++--
 .../discojs-core/src/weights/weights_container.ts |  3 ++-
 .../src/dataset/data_loader/image_loader.spec.ts  | 15 +++++++++------
 .../src/dataset/data_loader/image_loader.ts       |  6 ++++--
 .../dataset/data_loader/tabular_loader.spec.ts    |  3 ++-
 .../src/dataset/data_loader/tabular_loader.ts     |  6 ++++--
 discojs/discojs-node/src/index.ts                 |  1 -
 discojs/discojs-web/package.json                  |  3 ++-
 .../src/dataset/data_loader/image_loader.ts       |  4 +++-
 .../src/dataset/data_loader/tabular_loader.ts     |  4 +++-
 .../src/dataset/data_loader/text_loader.ts        |  4 +++-
 discojs/discojs-web/src/memory/memory.ts          |  3 ++-
 docs/CONTRIBUTING.md                              |  2 --
 package-lock.json                                 |  4 +++-
 server/package.json                               |  1 +
 server/src/get_server.ts                          |  4 +++-
 server/src/index.ts                               |  5 -----
 server/src/router/decentralized/server.ts         |  8 +++++---
 server/src/router/federated/server.ts             |  3 +--
 server/src/router/server.ts                       |  5 +++--
 server/src/router/tasks.ts                        | 11 +++++++----
 server/src/tasks.ts                               |  5 ++++-
 .../components/task_creation_form/TaskForm.vue    |  6 ++++--
 web-client/src/main.ts                            |  2 +-
 web-client/src/store/memory.ts                    |  4 +++-
 web-client/src/types.ts                           |  2 +-
 66 files changed, 215 insertions(+), 122 deletions(-)

diff --git a/cli/src/data.ts b/cli/src/data.ts
index 84baf71ca..c099ba9bf 100644
--- a/cli/src/data.ts
+++ b/cli/src/data.ts
@@ -3,7 +3,8 @@ import fs from 'node:fs'
 import fs_promises from 'fs/promises'
 import path from 'node:path'
 
-import { node, type data, type Task } from '@epfml/discojs-node'
+import type { Task, data } from '@epfml/discojs-node'
+import { node } from '@epfml/discojs-node'
 
 function filesFromFolder (dir: string, folder: string, fractionToKeep: number): string[] {
   const f = fs.readdirSync(dir + folder)
diff --git a/discojs/discojs-core/src/aggregator/base.ts b/discojs/discojs-core/src/aggregator/base.ts
index 7ceedd2fe..a9fe96c29 100644
--- a/discojs/discojs-core/src/aggregator/base.ts
+++ b/discojs/discojs-core/src/aggregator/base.ts
@@ -1,6 +1,7 @@
-import { type client, type Task, type tf, type AsyncInformant } from '..'
-
 import { List, Map, Set } from 'immutable'
+import type tf from '@tensorflow/tfjs'
+
+import type { client, Task, AsyncInformant } from '..'
 
 export enum AggregationStep {
   ADD,
diff --git a/discojs/discojs-core/src/aggregator/mean.spec.ts b/discojs/discojs-core/src/aggregator/mean.spec.ts
index 25f484e0e..5cbbc0da5 100644
--- a/discojs/discojs-core/src/aggregator/mean.spec.ts
+++ b/discojs/discojs-core/src/aggregator/mean.spec.ts
@@ -1,7 +1,9 @@
 import { assert, expect } from 'chai'
-import { type Map } from 'immutable'
+import type { Map } from 'immutable'
+import type tf from '@tensorflow/tfjs'
 
-import { aggregator, defaultTasks, type client, type Task, type tf } from '..'
+import type { client, Task } from '..'
+import { aggregator, defaultTasks } from '..'
 import { AggregationStep } from './base'
 
 const task = defaultTasks.titanic.getTask()
diff --git a/discojs/discojs-core/src/aggregator/mean.ts b/discojs/discojs-core/src/aggregator/mean.ts
index 5dcc24fdf..c964feab9 100644
--- a/discojs/discojs-core/src/aggregator/mean.ts
+++ b/discojs/discojs-core/src/aggregator/mean.ts
@@ -1,7 +1,9 @@
-import { type Map } from 'immutable'
+import type { Map } from 'immutable'
+import type tf from '@tensorflow/tfjs'
 
 import { AggregationStep, Base as Aggregator } from './base'
-import { type Task, type WeightsContainer, aggregation, type tf, type client } from '..'
+import type { Task, WeightsContainer, client } from '..'
+import { aggregation } from '..'
 
 /**
  * Mean aggregator whose aggregation step consists in computing the mean of the received weights.
diff --git a/discojs/discojs-core/src/aggregator/secure.ts b/discojs/discojs-core/src/aggregator/secure.ts
index 0a8c41eb4..a84111ea5 100644
--- a/discojs/discojs-core/src/aggregator/secure.ts
+++ b/discojs/discojs-core/src/aggregator/secure.ts
@@ -1,9 +1,10 @@
-import { AggregationStep, Base as Aggregator } from './base'
-import { tf, aggregation, type Task, type WeightsContainer, type client } from '..'
-
 import * as crypto from 'crypto'
-
 import { Map, List, Range } from 'immutable'
+import tf from '@tensorflow/tfjs'
+
+import { AggregationStep, Base as Aggregator } from './base'
+import type { Task, WeightsContainer, client } from '..'
+import { aggregation } from '..'
 
 /**
  * Aggregator implementing secure multi-party computation for decentralized learning.
diff --git a/discojs/discojs-core/src/client/base.ts b/discojs/discojs-core/src/client/base.ts
index 472dcd01c..821469141 100644
--- a/discojs/discojs-core/src/client/base.ts
+++ b/discojs/discojs-core/src/client/base.ts
@@ -1,10 +1,12 @@
 import axios from 'axios'
-import { type Set } from 'immutable'
+import type { Set } from 'immutable'
+import type tf from '@tensorflow/tfjs'
 
-import { type tf, type Task, type TrainingInformant, serialization, type WeightsContainer } from '..'
-import { type NodeID } from './types'
-import { type EventConnection } from './event_connection'
-import { type Aggregator } from '../aggregator'
+import type { Task, TrainingInformant, WeightsContainer } from '..'
+import { serialization } from '..'
+import type { NodeID } from './types'
+import type { EventConnection } from './event_connection'
+import type { Aggregator } from '../aggregator'
 
 /**
  * Main, abstract, class representing a Disco client in a network, which handles
diff --git a/discojs/discojs-core/src/dataset/data/data.ts b/discojs/discojs-core/src/dataset/data/data.ts
index b18462dfd..049b5a7f8 100644
--- a/discojs/discojs-core/src/dataset/data/data.ts
+++ b/discojs/discojs-core/src/dataset/data/data.ts
@@ -1,8 +1,9 @@
-import { type tf, type Task } from '../..'
-import { type Dataset } from '../dataset'
-import { type PreprocessingFunction } from './preprocessing/base'
+import type tf from '@tensorflow/tfjs'
+import type { List } from 'immutable'
 
-import { type List } from 'immutable'
+import type { Task } from '../..'
+import type { Dataset } from '../dataset'
+import type { PreprocessingFunction } from './preprocessing/base'
 
 /**
  * Abstract class representing an immutable Disco dataset, including a TF.js dataset,
diff --git a/discojs/discojs-core/src/dataset/data/image_data.spec.ts b/discojs/discojs-core/src/dataset/data/image_data.spec.ts
index 2d860c10b..c93120481 100644
--- a/discojs/discojs-core/src/dataset/data/image_data.spec.ts
+++ b/discojs/discojs-core/src/dataset/data/image_data.spec.ts
@@ -1,7 +1,8 @@
 import { assert, expect } from 'chai'
+import tf from '@tensorflow/tfjs'
 
 import { ImageData } from './image_data'
-import { tf, type Task } from '../..'
+import type { Task } from '../..'
 
 describe('image data checks', () => {
   const simplefaceMock: Task = {
diff --git a/discojs/discojs-core/src/dataset/data/image_data.ts b/discojs/discojs-core/src/dataset/data/image_data.ts
index 5fa097459..83a61117e 100644
--- a/discojs/discojs-core/src/dataset/data/image_data.ts
+++ b/discojs/discojs-core/src/dataset/data/image_data.ts
@@ -1,5 +1,7 @@
-import { type tf, type Task } from '../..'
-import { type Dataset } from '../dataset'
+import type tf from '@tensorflow/tfjs'
+
+import type { Task } from '../..'
+import type { Dataset } from '../dataset'
 import { Data } from './data'
 import { ImagePreprocessing, IMAGE_PREPROCESSING } from './preprocessing'
 
diff --git a/discojs/discojs-core/src/dataset/data/preprocessing/base.ts b/discojs/discojs-core/src/dataset/data/preprocessing/base.ts
index 8ab92f1d2..9507dfe6b 100644
--- a/discojs/discojs-core/src/dataset/data/preprocessing/base.ts
+++ b/discojs/discojs-core/src/dataset/data/preprocessing/base.ts
@@ -1,7 +1,9 @@
-import { type tf, type Task } from '../../..'
-import { type ImagePreprocessing } from './image_preprocessing'
-import { type TabularPreprocessing } from './tabular_preprocessing'
-import { type TextPreprocessing } from './text_preprocessing'
+import type tf from '@tensorflow/tfjs'
+
+import type { Task } from '../../..'
+import type { ImagePreprocessing } from './image_preprocessing'
+import type { TabularPreprocessing } from './tabular_preprocessing'
+import type { TextPreprocessing } from './text_preprocessing'
 
 /**
  * All available preprocessing type enums.
diff --git a/discojs/discojs-core/src/dataset/data/preprocessing/image_preprocessing.ts b/discojs/discojs-core/src/dataset/data/preprocessing/image_preprocessing.ts
index 9c8be655d..2220a9d69 100644
--- a/discojs/discojs-core/src/dataset/data/preprocessing/image_preprocessing.ts
+++ b/discojs/discojs-core/src/dataset/data/preprocessing/image_preprocessing.ts
@@ -1,7 +1,8 @@
-import { type Task, tf } from '../../..'
-import { type PreprocessingFunction } from './base'
-
 import { List } from 'immutable'
+import tf from '@tensorflow/tfjs'
+
+import type { Task } from '../../..'
+import type { PreprocessingFunction } from './base'
 
 /**
  * Available image preprocessing types.
diff --git a/discojs/discojs-core/src/dataset/data/preprocessing/tabular_preprocessing.ts b/discojs/discojs-core/src/dataset/data/preprocessing/tabular_preprocessing.ts
index 7cdc610d4..4450b41c2 100644
--- a/discojs/discojs-core/src/dataset/data/preprocessing/tabular_preprocessing.ts
+++ b/discojs/discojs-core/src/dataset/data/preprocessing/tabular_preprocessing.ts
@@ -1,6 +1,8 @@
-import { type Task, type tf } from '../../..'
+import type tf from '@tensorflow/tfjs'
 import { List } from 'immutable'
-import { type PreprocessingFunction } from './base'
+
+import type { Task } from '../../..'
+import type { PreprocessingFunction } from './base'
 
 /**
  * Available tabular preprocessing types.
diff --git a/discojs/discojs-core/src/dataset/data/preprocessing/text_preprocessing.ts b/discojs/discojs-core/src/dataset/data/preprocessing/text_preprocessing.ts
index 0ff7380f2..e372bde4d 100644
--- a/discojs/discojs-core/src/dataset/data/preprocessing/text_preprocessing.ts
+++ b/discojs/discojs-core/src/dataset/data/preprocessing/text_preprocessing.ts
@@ -1,7 +1,8 @@
-import { type Task, tf } from '../../..'
-import { type PreprocessingFunction } from './base'
-
 import { List } from 'immutable'
+import tf from '@tensorflow/tfjs'
+
+import type { Task } from '../../..'
+import type { PreprocessingFunction } from './base'
 
 /**
  * Available text preprocessing types.
diff --git a/discojs/discojs-core/src/dataset/data/tabular_data.spec.ts b/discojs/discojs-core/src/dataset/data/tabular_data.spec.ts
index 72362cc1b..26b610403 100644
--- a/discojs/discojs-core/src/dataset/data/tabular_data.spec.ts
+++ b/discojs/discojs-core/src/dataset/data/tabular_data.spec.ts
@@ -1,8 +1,9 @@
 import { assert, expect } from 'chai'
 import { Map, Set } from 'immutable'
+import tf from '@tensorflow/tfjs'
 
 import { TabularData } from './tabular_data'
-import { tf, type Task } from '../..'
+import type { Task } from '../..'
 
 describe('tabular data checks', () => {
   const titanicMock: Task = {
diff --git a/discojs/discojs-core/src/dataset/data_loader/image_loader.ts b/discojs/discojs-core/src/dataset/data_loader/image_loader.ts
index c3cb3cca4..aa3384c1c 100644
--- a/discojs/discojs-core/src/dataset/data_loader/image_loader.ts
+++ b/discojs/discojs-core/src/dataset/data_loader/image_loader.ts
@@ -1,9 +1,11 @@
 import { Range } from 'immutable'
+import tf from '@tensorflow/tfjs'
 
-import { tf } from '../..'
-import { type Dataset } from '../dataset'
-import { type Data, ImageData, type DataSplit } from '../data'
-import { DataLoader, type DataConfig } from '../data_loader'
+import type { Dataset } from '../dataset'
+import type { Data, DataSplit } from '../data'
+import { ImageData } from '../data'
+import type { DataConfig } from '../data_loader'
+import { DataLoader } from '../data_loader'
 
 /**
  * Image data loader whose instantiable implementation is delegated by the platform-dependent Disco subprojects, namely,
diff --git a/discojs/discojs-core/src/dataset/dataset.ts b/discojs/discojs-core/src/dataset/dataset.ts
index aaa825d13..121af728c 100644
--- a/discojs/discojs-core/src/dataset/dataset.ts
+++ b/discojs/discojs-core/src/dataset/dataset.ts
@@ -1,4 +1,4 @@
-import { type tf } from '..'
+import type tf from '@tensorflow/tfjs'
 
 /**
  * Convenient type for the common dataset type used in TF.js.
diff --git a/discojs/discojs-core/src/default_tasks/cifar10.ts b/discojs/discojs-core/src/default_tasks/cifar10.ts
index 32fdea95c..6b4cb478f 100644
--- a/discojs/discojs-core/src/default_tasks/cifar10.ts
+++ b/discojs/discojs-core/src/default_tasks/cifar10.ts
@@ -1,4 +1,7 @@
-import { tf, type Task, data, type TaskProvider } from '..'
+import tf from '@tensorflow/tfjs'
+
+import type { Task, TaskProvider } from '..'
+import { data } from '..'
 
 export const cifar10: TaskProvider = {
   getTask (): Task {
diff --git a/discojs/discojs-core/src/default_tasks/geotags.ts b/discojs/discojs-core/src/default_tasks/geotags.ts
index 149689054..06c594381 100644
--- a/discojs/discojs-core/src/default_tasks/geotags.ts
+++ b/discojs/discojs-core/src/default_tasks/geotags.ts
@@ -1,5 +1,8 @@
-import { tf, type Task, data, type TaskProvider } from '..'
 import { Range } from 'immutable'
+import tf from '@tensorflow/tfjs'
+
+import type { Task, TaskProvider } from '..'
+import { data } from '..'
 import { LabelTypeEnum } from '../task/label_type'
 
 export const geotags: TaskProvider = {
diff --git a/discojs/discojs-core/src/default_tasks/lus_covid.ts b/discojs/discojs-core/src/default_tasks/lus_covid.ts
index 2a1e27c47..f1b051956 100644
--- a/discojs/discojs-core/src/default_tasks/lus_covid.ts
+++ b/discojs/discojs-core/src/default_tasks/lus_covid.ts
@@ -1,4 +1,7 @@
-import { tf, type Task, data, type TaskProvider } from '..'
+import tf from '@tensorflow/tfjs'
+
+import type { Task, TaskProvider } from '..'
+import { data } from '..'
 
 export const lusCovid: TaskProvider = {
   getTask (): Task {
diff --git a/discojs/discojs-core/src/default_tasks/mnist.ts b/discojs/discojs-core/src/default_tasks/mnist.ts
index 627b045ef..d59421d86 100644
--- a/discojs/discojs-core/src/default_tasks/mnist.ts
+++ b/discojs/discojs-core/src/default_tasks/mnist.ts
@@ -1,4 +1,6 @@
-import { tf, type Task, type TaskProvider } from '..'
+import tf from '@tensorflow/tfjs'
+
+import type { Task, TaskProvider } from '..'
 
 export const mnist: TaskProvider = {
   getTask (): Task {
diff --git a/discojs/discojs-core/src/default_tasks/simple_face.ts b/discojs/discojs-core/src/default_tasks/simple_face.ts
index 20ef66ab7..d1d09a119 100644
--- a/discojs/discojs-core/src/default_tasks/simple_face.ts
+++ b/discojs/discojs-core/src/default_tasks/simple_face.ts
@@ -1,4 +1,7 @@
-import { type tf, type Task, data, type TaskProvider } from '..'
+import type tf from '@tensorflow/tfjs'
+
+import type { Task, TaskProvider } from '..'
+import { data } from '..'
 
 export const simpleFace: TaskProvider = {
   getTask (): Task {
diff --git a/discojs/discojs-core/src/default_tasks/skin_mnist.ts b/discojs/discojs-core/src/default_tasks/skin_mnist.ts
index 81ec2b363..4b41bb581 100644
--- a/discojs/discojs-core/src/default_tasks/skin_mnist.ts
+++ b/discojs/discojs-core/src/default_tasks/skin_mnist.ts
@@ -1,5 +1,7 @@
+import tf from '@tensorflow/tfjs'
+
 import type { Task, TaskProvider } from '..'
-import { tf, data } from '..'
+import { data } from '..'
 
 export const skinMnist: TaskProvider = {
   getTask (): Task {
diff --git a/discojs/discojs-core/src/default_tasks/titanic.ts b/discojs/discojs-core/src/default_tasks/titanic.ts
index 5157c1f3e..ea231cb67 100644
--- a/discojs/discojs-core/src/default_tasks/titanic.ts
+++ b/discojs/discojs-core/src/default_tasks/titanic.ts
@@ -1,4 +1,7 @@
-import { tf, type Task, type TaskProvider, data } from '..'
+import tf from '@tensorflow/tfjs'
+
+import type { Task, TaskProvider } from '..'
+import { data } from '..'
 
 export const titanic: TaskProvider = {
   getTask (): Task {
diff --git a/discojs/discojs-core/src/index.ts b/discojs/discojs-core/src/index.ts
index d2693106a..4e19e8d2b 100644
--- a/discojs/discojs-core/src/index.ts
+++ b/discojs/discojs-core/src/index.ts
@@ -1,5 +1,3 @@
-export * as tf from '@tensorflow/tfjs'
-
 export * as data from './dataset'
 export * as serialization from './serialization'
 export * as training from './training'
diff --git a/discojs/discojs-core/src/logging/trainer_logger.ts b/discojs/discojs-core/src/logging/trainer_logger.ts
index 1822ffd39..84bcdec6a 100644
--- a/discojs/discojs-core/src/logging/trainer_logger.ts
+++ b/discojs/discojs-core/src/logging/trainer_logger.ts
@@ -1,6 +1,6 @@
 import { List } from 'immutable'
+import tf from '@tensorflow/tfjs'
 
-import { tf } from '..'
 import { ConsoleLogger } from '.'
 
 export class TrainerLog {
diff --git a/discojs/discojs-core/src/memory/base.ts b/discojs/discojs-core/src/memory/base.ts
index 43ff5fc7a..59b4c85a8 100644
--- a/discojs/discojs-core/src/memory/base.ts
+++ b/discojs/discojs-core/src/memory/base.ts
@@ -1,9 +1,9 @@
 // only used browser-side
 // TODO: replace IO type
-import type * as tf from '@tensorflow/tfjs'
+import type tf from '@tensorflow/tfjs'
 
-import { type TaskID } from '..'
-import { type ModelType } from './model_type'
+import type { TaskID } from '..'
+import type { ModelType } from './model_type'
 
 /**
  * Model path which uniquely identifies a model in memory.
diff --git a/discojs/discojs-core/src/memory/empty.ts b/discojs/discojs-core/src/memory/empty.ts
index 875964753..6abeb2536 100644
--- a/discojs/discojs-core/src/memory/empty.ts
+++ b/discojs/discojs-core/src/memory/empty.ts
@@ -1,6 +1,7 @@
-import { type tf } from '..'
+import type tf from '@tensorflow/tfjs'
 
-import { Memory, type ModelInfo, type Path } from './base'
+import type { ModelInfo, Path } from './base'
+import { Memory } from './base'
 
 /**
  * Represents an empty model memory.
diff --git a/discojs/discojs-core/src/privacy.ts b/discojs/discojs-core/src/privacy.ts
index 4eb5021e8..754410a6b 100644
--- a/discojs/discojs-core/src/privacy.ts
+++ b/discojs/discojs-core/src/privacy.ts
@@ -1,4 +1,6 @@
-import { tf, type Task, type WeightsContainer } from '.'
+import tf from '@tensorflow/tfjs'
+
+import type { Task, WeightsContainer } from '.'
 
 /**
  * Add task-parametrized Gaussian noise to and clip the weights update between the previous and current rounds.
diff --git a/discojs/discojs-core/src/serialization/model.spec.ts b/discojs/discojs-core/src/serialization/model.spec.ts
index b544501d5..728aac4ba 100644
--- a/discojs/discojs-core/src/serialization/model.spec.ts
+++ b/discojs/discojs-core/src/serialization/model.spec.ts
@@ -1,6 +1,7 @@
 import { assert } from 'chai'
+import tf from '@tensorflow/tfjs'
 
-import { tf, serialization } from '..'
+import { serialization } from '..'
 
 async function getRawWeights (model: tf.LayersModel): Promise<Array<[number, Float32Array]>> {
   return Array.from(
diff --git a/discojs/discojs-core/src/serialization/model.ts b/discojs/discojs-core/src/serialization/model.ts
index 668685ac7..ac8776569 100644
--- a/discojs/discojs-core/src/serialization/model.ts
+++ b/discojs/discojs-core/src/serialization/model.ts
@@ -1,4 +1,4 @@
-import { tf } from '..'
+import tf from '@tensorflow/tfjs'
 import msgpack from 'msgpack-lite'
 
 export type Encoded = number[]
diff --git a/discojs/discojs-core/src/serialization/weights.ts b/discojs/discojs-core/src/serialization/weights.ts
index 831f69ea8..59d687a55 100644
--- a/discojs/discojs-core/src/serialization/weights.ts
+++ b/discojs/discojs-core/src/serialization/weights.ts
@@ -1,6 +1,7 @@
 import * as msgpack from 'msgpack-lite'
+import tf from '@tensorflow/tfjs'
 
-import { tf, WeightsContainer } from '..'
+import { WeightsContainer } from '..'
 
 interface Serialized {
   shape: number[]
diff --git a/discojs/discojs-core/src/task/task_handler.ts b/discojs/discojs-core/src/task/task_handler.ts
index 8a5540759..9a70f5e1b 100644
--- a/discojs/discojs-core/src/task/task_handler.ts
+++ b/discojs/discojs-core/src/task/task_handler.ts
@@ -1,8 +1,10 @@
 import axios from 'axios'
 import { Map } from 'immutable'
+import type tf from '@tensorflow/tfjs'
 
-import { serialization, type tf, WeightsContainer } from '..'
-import { isTask, type Task, type TaskID } from './task'
+import { serialization, WeightsContainer } from '..'
+import type { Task, TaskID } from './task'
+import { isTask } from './task'
 
 const TASK_ENDPOINT = 'tasks'
 
diff --git a/discojs/discojs-core/src/task/task_provider.ts b/discojs/discojs-core/src/task/task_provider.ts
index bb4c3b586..31af016bd 100644
--- a/discojs/discojs-core/src/task/task_provider.ts
+++ b/discojs/discojs-core/src/task/task_provider.ts
@@ -1,4 +1,6 @@
-import { type tf, type Task } from '..'
+import type tf from '@tensorflow/tfjs'
+
+import type { Task } from '..'
 
 export interface TaskProvider {
   getTask: () => Task
diff --git a/discojs/discojs-core/src/training/trainer/distributed_trainer.ts b/discojs/discojs-core/src/training/trainer/distributed_trainer.ts
index c54ff8776..61bf94632 100644
--- a/discojs/discojs-core/src/training/trainer/distributed_trainer.ts
+++ b/discojs/discojs-core/src/training/trainer/distributed_trainer.ts
@@ -1,5 +1,8 @@
-import { type tf, type Memory, type Task, type TrainingInformant, type TrainingFunction, WeightsContainer, type client as clients } from '../..'
-import { type Aggregator } from '../../aggregator'
+import type tf from '@tensorflow/tfjs'
+
+import type { Memory, Task, TrainingInformant, TrainingFunction, client as clients } from '../..'
+import { WeightsContainer } from '../..'
+import type { Aggregator } from '../../aggregator'
 import { Trainer } from './trainer'
 
 /**
diff --git a/discojs/discojs-core/src/training/trainer/local_trainer.ts b/discojs/discojs-core/src/training/trainer/local_trainer.ts
index aefe1adc5..521982a9e 100644
--- a/discojs/discojs-core/src/training/trainer/local_trainer.ts
+++ b/discojs/discojs-core/src/training/trainer/local_trainer.ts
@@ -1,4 +1,5 @@
-import { type tf } from '../..'
+import type tf from '@tensorflow/tfjs'
+
 import { Trainer } from './trainer'
 
 /** Class whose role is to locally (alone) train a model on a given dataset,
diff --git a/discojs/discojs-core/src/training/trainer/trainer.ts b/discojs/discojs-core/src/training/trainer/trainer.ts
index 0b6980e80..fea2a8dcf 100644
--- a/discojs/discojs-core/src/training/trainer/trainer.ts
+++ b/discojs/discojs-core/src/training/trainer/trainer.ts
@@ -1,7 +1,11 @@
-import { type tf, type Memory, type Task, type TrainingInformant, type TrainingFunction, fitModelFunctions } from '../..'
+import type tf from '@tensorflow/tfjs'
+
+import type { Memory, Task, TrainingInformant, TrainingFunction } from '../..'
+import { fitModelFunctions } from '../..'
 
 import { RoundTracker } from './round_tracker'
-import { TrainerLogger, type TrainerLog } from '../../logging/trainer_logger'
+import type { TrainerLog } from '../../logging/trainer_logger'
+import { TrainerLogger } from '../../logging/trainer_logger'
 
 /** Abstract class whose role is to train a model with a given dataset. This can be either done
  * locally (alone) or in a distributed way with collaborators. The Trainer works as follows:
diff --git a/discojs/discojs-core/src/training/trainer/trainer_builder.ts b/discojs/discojs-core/src/training/trainer/trainer_builder.ts
index 4b52562b0..c32122a43 100644
--- a/discojs/discojs-core/src/training/trainer/trainer_builder.ts
+++ b/discojs/discojs-core/src/training/trainer/trainer_builder.ts
@@ -1,9 +1,12 @@
-import { type tf, type client as clients, type Task, type TrainingInformant, type TrainingFunction, type Memory, ModelType, type ModelInfo } from '../..'
-import { type Aggregator } from '../../aggregator'
+import type tf from '@tensorflow/tfjs'
+
+import type { client as clients, Task, TrainingInformant, TrainingFunction, ModelInfo, Memory } from '../..'
+import { ModelType } from '../..'
+import type { Aggregator } from '../../aggregator'
 
 import { DistributedTrainer } from './distributed_trainer'
 import { LocalTrainer } from './local_trainer'
-import { type Trainer } from './trainer'
+import type { Trainer } from './trainer'
 
 /**
  * A class that helps build the Trainer and auxiliary classes.
diff --git a/discojs/discojs-core/src/training/training_functions.ts b/discojs/discojs-core/src/training/training_functions.ts
index 861e03cdd..286cc9622 100644
--- a/discojs/discojs-core/src/training/training_functions.ts
+++ b/discojs/discojs-core/src/training/training_functions.ts
@@ -1,4 +1,6 @@
-import { type tf, type TrainingInformation } from '..'
+import type tf from '@tensorflow/tfjs'
+
+import type { TrainingInformation } from '..'
 
 /**
  * The training function is the function that will be called to train the model.
diff --git a/discojs/discojs-core/src/types.ts b/discojs/discojs-core/src/types.ts
index ab71494a9..84fc9ff68 100644
--- a/discojs/discojs-core/src/types.ts
+++ b/discojs/discojs-core/src/types.ts
@@ -1,6 +1,9 @@
-import { type WeightsContainer, type tf } from '.'
-import { type NodeID } from './client'
-import { type Map } from 'immutable'
+import type tf from '@tensorflow/tfjs'
+import type { Map } from 'immutable'
+
+import type { WeightsContainer } from '.'
+import type { NodeID } from './client'
+
 // Filesystem reference
 export type Path = string
 
diff --git a/discojs/discojs-core/src/validation/validator.ts b/discojs/discojs-core/src/validation/validator.ts
index 59ba5c249..832613a3a 100644
--- a/discojs/discojs-core/src/validation/validator.ts
+++ b/discojs/discojs-core/src/validation/validator.ts
@@ -1,6 +1,8 @@
 import { List } from 'immutable'
+import tf from '@tensorflow/tfjs'
 
-import { tf, type data, type Task, type Logger, type client as clients, GraphInformant, type Memory, type ModelSource, type Features } from '..'
+import type { data, Task, Logger, client as clients, Memory, ModelSource, Features } from '..'
+import { GraphInformant } from '..'
 
 export class Validator {
   private readonly graphInformant = new GraphInformant()
diff --git a/discojs/discojs-core/src/weights/aggregation.ts b/discojs/discojs-core/src/weights/aggregation.ts
index fd694c452..b3973b47f 100644
--- a/discojs/discojs-core/src/weights/aggregation.ts
+++ b/discojs/discojs-core/src/weights/aggregation.ts
@@ -1,7 +1,8 @@
 import { List } from 'immutable'
+import tf from '@tensorflow/tfjs'
 
-import { tf } from '..'
-import { type TensorLike, WeightsContainer } from './weights_container'
+import type { TensorLike } from './weights_container'
+import { WeightsContainer } from './weights_container'
 
 type WeightsLike = Iterable<TensorLike>
 
diff --git a/discojs/discojs-core/src/weights/weights_container.ts b/discojs/discojs-core/src/weights/weights_container.ts
index f12070c82..6bc262844 100644
--- a/discojs/discojs-core/src/weights/weights_container.ts
+++ b/discojs/discojs-core/src/weights/weights_container.ts
@@ -1,6 +1,7 @@
 import { List } from 'immutable'
+import tf from '@tensorflow/tfjs'
 
-import { tf, type Weights } from '..'
+import type { Weights } from '..'
 
 export type TensorLike = tf.Tensor | ArrayLike<number>
 
diff --git a/discojs/discojs-node/src/dataset/data_loader/image_loader.spec.ts b/discojs/discojs-node/src/dataset/data_loader/image_loader.spec.ts
index 291cb7bb1..c38417b56 100644
--- a/discojs/discojs-node/src/dataset/data_loader/image_loader.spec.ts
+++ b/discojs/discojs-node/src/dataset/data_loader/image_loader.spec.ts
@@ -2,7 +2,10 @@ import { assert, expect } from 'chai'
 import { List, Map, Range } from 'immutable'
 import fs from 'fs'
 
-import { tf, node, type Task } from '../..'
+import tf from '@tensorflow/tfjs'
+import { node as tfNode } from '@tensorflow/tfjs-node'
+
+import { node, type Task } from '../..'
 
 const readFilesFromDir = (dir: string): string[] =>
   fs.readdirSync(dir).map((file: string) => dir + file)
@@ -41,14 +44,14 @@ describe('image loader', () => {
   it('loads single sample without label', async () => {
     const file = '../../example_training_data/9-mnist-example.png'
     const singletonDataset = await LOADERS.MNIST.load(file)
-    const imageContent = tf.node.decodeImage(fs.readFileSync(file))
+    const imageContent = tfNode.decodeImage(fs.readFileSync(file))
     await Promise.all((await singletonDataset.toArrayForTest()).map(async (entry) => {
       expect(await imageContent.bytes()).eql(await (entry as tf.Tensor).bytes())
     }))
   })
 
   it('loads multiple samples without labels', async () => {
-    const imagesContent = FILES.CIFAR10.map((file) => tf.node.decodeImage(fs.readFileSync(file)))
+    const imagesContent = FILES.CIFAR10.map((file) => tfNode.decodeImage(fs.readFileSync(file)))
     const datasetContent = await (await LOADERS.CIFAR10
       .loadAll(FILES.CIFAR10, { shuffle: false }))
       .train.dataset.toArray()
@@ -58,7 +61,7 @@ describe('image loader', () => {
 
   it('loads single sample with label', async () => {
     const path = DIRS.CIFAR10 + '0.png'
-    const imageContent = tf.node.decodeImage(fs.readFileSync(path))
+    const imageContent = tfNode.decodeImage(fs.readFileSync(path))
     const datasetContent = await (await LOADERS.CIFAR10
       .load(path, { labels: ['example'] })).toArray()
     expect((datasetContent[0] as any).xs.shape).eql(imageContent.shape)
@@ -70,7 +73,7 @@ describe('image loader', () => {
     const stringLabels = labels.map((label) => label.toString())
     const oneHotLabels = List(tf.oneHot(labels.toArray(), 10).arraySync() as number[])
 
-    const imagesContent = List(FILES.CIFAR10.map((file) => tf.node.decodeImage(fs.readFileSync(file))))
+    const imagesContent = List(FILES.CIFAR10.map((file) => tfNode.decodeImage(fs.readFileSync(file))))
     const datasetContent = List(await (await LOADERS.CIFAR10
       .loadAll(FILES.CIFAR10, { labels: stringLabels.toArray(), shuffle: false }))
       .train.dataset.toArray())
@@ -125,7 +128,7 @@ describe('image loader', () => {
   })
   it('validation split', async () => {
     const validationSplit = 0.2
-    const imagesContent = FILES.CIFAR10.map((file) => tf.node.decodeImage(fs.readFileSync(file)))
+    const imagesContent = FILES.CIFAR10.map((file) => tfNode.decodeImage(fs.readFileSync(file)))
     const datasetContent = await new node.data.NodeImageLoader(cifar10Mock)
       .loadAll(FILES.CIFAR10, { shuffle: false, validationSplit })
 
diff --git a/discojs/discojs-node/src/dataset/data_loader/image_loader.ts b/discojs/discojs-node/src/dataset/data_loader/image_loader.ts
index eccec2954..58b1d7954 100644
--- a/discojs/discojs-node/src/dataset/data_loader/image_loader.ts
+++ b/discojs/discojs-node/src/dataset/data_loader/image_loader.ts
@@ -1,9 +1,11 @@
 import fs from 'fs'
+import type tf from '@tensorflow/tfjs'
+import { node } from '@tensorflow/tfjs-node'
 
-import { tf, data } from '../..'
+import { data } from '@epfml/discojs-core'
 
 export class NodeImageLoader extends data.ImageLoader<string> {
   async readImageFrom (source: string): Promise<tf.Tensor3D> {
-    return tf.node.decodeImage(fs.readFileSync(source)) as tf.Tensor3D
+    return node.decodeImage(fs.readFileSync(source)) as tf.Tensor3D
   }
 }
diff --git a/discojs/discojs-node/src/dataset/data_loader/tabular_loader.spec.ts b/discojs/discojs-node/src/dataset/data_loader/tabular_loader.spec.ts
index 6b4191737..d1745556e 100644
--- a/discojs/discojs-node/src/dataset/data_loader/tabular_loader.spec.ts
+++ b/discojs/discojs-node/src/dataset/data_loader/tabular_loader.spec.ts
@@ -1,7 +1,8 @@
 import { assert, expect } from 'chai'
 import { List } from 'immutable'
+import tf from '@tensorflow/tfjs'
 
-import { tf, node, type Task } from '../..'
+import { node, type Task } from '../..'
 
 const inputFiles = ['../../example_training_data/titanic_train.csv']
 
diff --git a/discojs/discojs-node/src/dataset/data_loader/tabular_loader.ts b/discojs/discojs-node/src/dataset/data_loader/tabular_loader.ts
index 8da13ea56..79aec87a2 100644
--- a/discojs/discojs-node/src/dataset/data_loader/tabular_loader.ts
+++ b/discojs/discojs-node/src/dataset/data_loader/tabular_loader.ts
@@ -1,4 +1,6 @@
-import { tf, data } from '../..'
+import { data as tfData } from '@tensorflow/tfjs-node'
+
+import { data } from '@epfml/discojs-core'
 
 export class NodeTabularLoader extends data.TabularLoader<string> {
   async loadDatasetFrom (source: string, csvConfig: Record<string, unknown>): Promise<data.Dataset> {
@@ -6,6 +8,6 @@ export class NodeTabularLoader extends data.TabularLoader<string> {
     if (source.slice(0, 7) !== prefix) {
       source = prefix + source
     }
-    return tf.data.csv(source, csvConfig)
+    return tfData.csv(source, csvConfig)
   }
 }
diff --git a/discojs/discojs-node/src/index.ts b/discojs/discojs-node/src/index.ts
index c2457af94..8b4108de4 100644
--- a/discojs/discojs-node/src/index.ts
+++ b/discojs/discojs-node/src/index.ts
@@ -1,4 +1,3 @@
 export * from '@epfml/discojs-core'
-export * as tf from '@tensorflow/tfjs-node'
 
 export * as node from './imports'
diff --git a/discojs/discojs-web/package.json b/discojs/discojs-web/package.json
index ba1da9f9b..6fe759f07 100644
--- a/discojs/discojs-web/package.json
+++ b/discojs/discojs-web/package.json
@@ -19,7 +19,8 @@
   },
   "homepage": "https://github.com/epfml/disco#readme",
   "dependencies": {
-    "@epfml/discojs-core": "*"
+    "@epfml/discojs-core": "*",
+    "@tensorflow/tfjs": "4"
   },
   "devDependencies": {
     "nodemon": "3"
diff --git a/discojs/discojs-web/src/dataset/data_loader/image_loader.ts b/discojs/discojs-web/src/dataset/data_loader/image_loader.ts
index d7116c9fb..6cabcc0b7 100644
--- a/discojs/discojs-web/src/dataset/data_loader/image_loader.ts
+++ b/discojs/discojs-web/src/dataset/data_loader/image_loader.ts
@@ -1,4 +1,6 @@
-import { tf, data } from '../..'
+import tf from '@tensorflow/tfjs'
+
+import { data } from '../..'
 
 export class WebImageLoader extends data.ImageLoader<File> {
   async readImageFrom (source: File): Promise<tf.Tensor3D> {
diff --git a/discojs/discojs-web/src/dataset/data_loader/tabular_loader.ts b/discojs/discojs-web/src/dataset/data_loader/tabular_loader.ts
index 8743defbd..f4ab7425c 100644
--- a/discojs/discojs-web/src/dataset/data_loader/tabular_loader.ts
+++ b/discojs/discojs-web/src/dataset/data_loader/tabular_loader.ts
@@ -1,4 +1,6 @@
-import { tf, data } from '../..'
+import tf from '@tensorflow/tfjs'
+
+import { data } from '../..'
 
 export class WebTabularLoader extends data.TabularLoader<File> {
   async loadDatasetFrom (source: File, csvConfig: Record<string, unknown>): Promise<data.Dataset> {
diff --git a/discojs/discojs-web/src/dataset/data_loader/text_loader.ts b/discojs/discojs-web/src/dataset/data_loader/text_loader.ts
index 6662ec734..eab2c98b5 100644
--- a/discojs/discojs-web/src/dataset/data_loader/text_loader.ts
+++ b/discojs/discojs-web/src/dataset/data_loader/text_loader.ts
@@ -1,4 +1,6 @@
-import { tf, data } from '../..'
+import tf from '@tensorflow/tfjs'
+
+import { data } from '../..'
 
 export class WebTextLoader extends data.TextLoader<File> {
   async loadDatasetFrom (source: File, config?: Record<string, unknown>): Promise<data.Dataset> {
diff --git a/discojs/discojs-web/src/memory/memory.ts b/discojs/discojs-web/src/memory/memory.ts
index 721af5863..d821812c7 100644
--- a/discojs/discojs-web/src/memory/memory.ts
+++ b/discojs/discojs-web/src/memory/memory.ts
@@ -8,8 +8,9 @@
  */
 import { Map } from 'immutable'
 import path from 'path'
+import * as tf from '@tensorflow/tfjs'
 
-import { tf, Memory, ModelType, type Path, type ModelInfo, type ModelSource } from '..'
+import { Memory, ModelType, type Path, type ModelInfo, type ModelSource } from '..'
 
 export class IndexedDB extends Memory {
   pathFor (source: ModelSource): Path {
diff --git a/docs/CONTRIBUTING.md b/docs/CONTRIBUTING.md
index 6dc2a4fe8..0288ec39c 100644
--- a/docs/CONTRIBUTING.md
+++ b/docs/CONTRIBUTING.md
@@ -142,8 +142,6 @@ Similarly to the server, any file ending with `.spec.ts` will be ran in the test
 
 `discojs-core` contains the core, platform-agnostic code of Disco.js, used by both `discojs-web` and `discojs-node`. As such, contributions to `discojs-core` must only contain code independent of either Node or the browser. As the names subtly suggest, `discojs-node` and `discojs-web` implement features specific to Node.js and browsers respectively, mostly related to memory and data handling as browser don't allow access to the file system.
 
-Note that, if you end up making calls to the Tensorflow.js API, you must import it from the root index. This is to ensure the right version of TF.js is loaded (depending on the compilation `dist/`), and only once. The only exception occurs in unit tests, which should import TF.js from the (local) `@epfml/discojs-node`, since those run on Node.js.
-
 Currently, the `discojs-node` project is available as the `@epfml/discojs-node` NPM package, which can be installed with
 `npm i @epfml/discojs-node` while the `discojs-web` project is available as the `@epfml/discojs` (and **not** as `@epfml/discojs-web`) NPM package, which can be installed with `npm i @epfml/discojs`.
 
diff --git a/package-lock.json b/package-lock.json
index 89baf9c9c..571eb3839 100644
--- a/package-lock.json
+++ b/package-lock.json
@@ -84,7 +84,8 @@
       "name": "@epfml/discojs",
       "version": "2.1.1",
       "dependencies": {
-        "@epfml/discojs-core": "*"
+        "@epfml/discojs-core": "*",
+        "@tensorflow/tfjs": "4"
       },
       "devDependencies": {
         "nodemon": "3"
@@ -20128,6 +20129,7 @@
       "dependencies": {
         "@epfml/discojs-node": "*",
         "@koush/wrtc": "0.5",
+        "@tensorflow/tfjs": "4",
         "cors": "2",
         "express": "4",
         "express-ws": "5",
diff --git a/server/package.json b/server/package.json
index 9f7d37a23..d4f68eb56 100644
--- a/server/package.json
+++ b/server/package.json
@@ -17,6 +17,7 @@
   "dependencies": {
     "@epfml/discojs-node": "*",
     "@koush/wrtc": "0.5",
+    "@tensorflow/tfjs": "4",
     "cors": "2",
     "express": "4",
     "express-ws": "5",
diff --git a/server/src/get_server.ts b/server/src/get_server.ts
index ae752e8c8..a124f77f8 100644
--- a/server/src/get_server.ts
+++ b/server/src/get_server.ts
@@ -1,11 +1,13 @@
 import cors from 'cors'
 import express from 'express'
 import expressWS from 'express-ws'
+import type tf from '@tensorflow/tfjs'
+
+import type { Task, TaskProvider } from '@epfml/discojs-node'
 
 import { CONFIG } from './config'
 import { Router } from './router'
 import { TasksAndModels } from './tasks'
-import { type tf, type Task, type TaskProvider } from '@epfml/discojs-node'
 import type * as http from 'http'
 
 export class Disco {
diff --git a/server/src/index.ts b/server/src/index.ts
index 9303830a4..b30c01adc 100644
--- a/server/src/index.ts
+++ b/server/src/index.ts
@@ -1,6 +1 @@
-import { tf } from '@epfml/discojs-node'
-
 export { Disco } from './get_server'
-
-// TODO, can we let the user of the server retrieve its own tfjs to get the types?
-export { tf }
diff --git a/server/src/router/decentralized/server.ts b/server/src/router/decentralized/server.ts
index 4713ade32..d779720a9 100644
--- a/server/src/router/decentralized/server.ts
+++ b/server/src/router/decentralized/server.ts
@@ -2,11 +2,13 @@ import { v4 as randomUUID } from 'uuid'
 import type express from 'express'
 import msgpack from 'msgpack-lite'
 import type WebSocket from 'ws'
-import { type ParamsDictionary } from 'express-serve-static-core'
-import { type ParsedQs } from 'qs'
+import type { ParamsDictionary } from 'express-serve-static-core'
+import type { ParsedQs } from 'qs'
 import { Map, Set } from 'immutable'
+import type tf from '@tensorflow/tfjs'
 
-import { type tf, client, type Task, type TaskID } from '@epfml/discojs-node'
+import type { Task, TaskID } from '@epfml/discojs-node'
+import { client } from '@epfml/discojs-node'
 
 import { Server } from '../server'
 
diff --git a/server/src/router/federated/server.ts b/server/src/router/federated/server.ts
index e6e141d1c..5444a15b1 100644
--- a/server/src/router/federated/server.ts
+++ b/server/src/router/federated/server.ts
@@ -1,13 +1,12 @@
 import type express from 'express'
 import type WebSocket from 'ws'
 import { v4 as randomUUID } from 'uuid'
-
 import { List, Map } from 'immutable'
 import msgpack from 'msgpack-lite'
+import type tf from '@tensorflow/tfjs'
 
 import {
   client,
-  type tf,
   serialization,
   AsyncInformant,
   type Task,
diff --git a/server/src/router/server.ts b/server/src/router/server.ts
index 507bc85f6..7ee7ba9fb 100644
--- a/server/src/router/server.ts
+++ b/server/src/router/server.ts
@@ -1,10 +1,11 @@
 import express from 'express'
 import type expressWS from 'express-ws'
 import type WebSocket from 'ws'
+import type tf from '@tensorflow/tfjs'
 
-import { type tf, type Task } from '@epfml/discojs-node'
+import type { Task } from '@epfml/discojs-node'
 
-import { type TasksAndModels } from '../tasks'
+import type { TasksAndModels } from '../tasks'
 
 export abstract class Server {
   private readonly ownRouter: expressWS.Router
diff --git a/server/src/router/tasks.ts b/server/src/router/tasks.ts
index 4e5af8b68..51bc58fe3 100644
--- a/server/src/router/tasks.ts
+++ b/server/src/router/tasks.ts
@@ -1,10 +1,13 @@
-import express, { type Request, type Response } from 'express'
+import type { Request, Response } from 'express'
+import express from 'express'
 import { Set } from 'immutable'
+import type tf from '@tensorflow/tfjs'
 
-import { type tf, serialization, type Task, type TaskID, isTask } from '@epfml/discojs-node'
+import type { Task, TaskID } from '@epfml/discojs-node'
+import { serialization, isTask } from '@epfml/discojs-node'
 
-import { type Config } from '../config'
-import { type TasksAndModels } from '../tasks'
+import type { Config } from '../config'
+import type { TasksAndModels } from '../tasks'
 
 export class Tasks {
   private readonly ownRouter: express.Router
diff --git a/server/src/tasks.ts b/server/src/tasks.ts
index c792c8880..25f23b7f0 100644
--- a/server/src/tasks.ts
+++ b/server/src/tasks.ts
@@ -1,8 +1,11 @@
 import { List, Set } from 'immutable'
 import { createHash } from 'node:crypto'
 import fs from 'node:fs'
+import tf from '@tensorflow/tfjs'
+import '@tensorflow/tfjs-node'
 
-import { tf, type Task, type Path, type Digest, isTaskProvider, type TaskProvider, defaultTasks } from '@epfml/discojs-node'
+import type { Task, Path, Digest, TaskProvider } from '@epfml/discojs-node'
+import { isTaskProvider, defaultTasks } from '@epfml/discojs-node'
 
 // default tasks and added ones
 // register 'taskAndModel' event to get tasks
diff --git a/web-client/src/components/task_creation_form/TaskForm.vue b/web-client/src/components/task_creation_form/TaskForm.vue
index b07e16c36..be92194d8 100644
--- a/web-client/src/components/task_creation_form/TaskForm.vue
+++ b/web-client/src/components/task_creation_form/TaskForm.vue
@@ -153,10 +153,12 @@
 <script lang="ts" setup>
 import * as yup from 'yup'
 import { ref, shallowRef } from 'vue'
-import { List, Map } from 'immutable'
 import { Form as VeeForm, ErrorMessage } from 'vee-validate'
 
-import { tf, Task, pushTask } from '@epfml/discojs'
+import { List, Map } from 'immutable'
+import * as tf from '@tensorflow/tfjs'
+
+import { Task, pushTask } from '@epfml/discojs'
 
 import { sections, modelCompileData, privacyParameters, FormField, FormSection } from '@/task_creation_form'
 import { useToaster } from '@/composables/toaster'
diff --git a/web-client/src/main.ts b/web-client/src/main.ts
index 3d936696e..9f7008eb5 100644
--- a/web-client/src/main.ts
+++ b/web-client/src/main.ts
@@ -10,7 +10,7 @@ import { createCustomI18n } from './locales/i18n'
 import '@/assets/css/tailwind.css'
 import '@/assets/css/styles.css'
 
-import { tf } from '@epfml/discojs'
+import * as tf from '@tensorflow/tfjs'
 
 tf.ready()
   .then(() => console.log(`Loaded ${tf.getBackend()} backend`))
diff --git a/web-client/src/store/memory.ts b/web-client/src/store/memory.ts
index 892787606..0dd04b7bc 100644
--- a/web-client/src/store/memory.ts
+++ b/web-client/src/store/memory.ts
@@ -1,8 +1,10 @@
 import { defineStore } from 'pinia'
 import { ref, shallowRef } from 'vue'
+
 import { Map } from 'immutable'
+import * as tf from '@tensorflow/tfjs'
 
-import { ModelInfo, ModelType, Path, tf } from '@epfml/discojs'
+import { ModelInfo, ModelType, Path } from '@epfml/discojs'
 
 interface ModelMetadata extends ModelInfo {
   date: string
diff --git a/web-client/src/types.ts b/web-client/src/types.ts
index c684084fc..38d3e043a 100644
--- a/web-client/src/types.ts
+++ b/web-client/src/types.ts
@@ -1,4 +1,4 @@
-import { tf } from '@epfml/discojs'
+import type tf from '@tensorflow/tfjs'
 
 export type Weights = tf.Tensor[]
 

From 9253c4713dd84d2877e63a9756a24ad65a10a209 Mon Sep 17 00:00:00 2001
From: tharvik <tharvik@users.noreply.github.com>
Date: Wed, 28 Feb 2024 19:13:10 +0100
Subject: [PATCH 2/2] discojs-{node,web}: don't reexport core

Closes: #638
---
 cli/src/args.ts                                 |  3 ++-
 cli/src/cli.ts                                  |  3 ++-
 cli/src/data.ts                                 | 10 +++++-----
 cli/src/utils.ts                                |  2 +-
 .../discojs-core/src/aggregator/secure.spec.ts  |  2 +-
 .../src/validation/validator.spec.ts            | 11 +++++------
 .../src/weights/aggregation.spec.ts             |  2 +-
 .../data_loader => data}/image_loader.spec.ts   | 17 +++++++++--------
 .../data_loader => data}/image_loader.ts        |  2 +-
 discojs/discojs-node/src/data/index.ts          |  2 ++
 .../data_loader => data}/tabular_loader.spec.ts |  8 +++++---
 .../data_loader => data}/tabular_loader.ts      |  2 +-
 .../data_loader => data}/text_loader.ts         |  0
 .../src/dataset/data_loader/index.ts            |  2 --
 discojs/discojs-node/src/imports.ts             |  1 -
 discojs/discojs-node/src/index.ts               |  4 +---
 .../data_loader => data}/image_loader.ts        |  4 ++--
 discojs/discojs-web/src/data/index.ts           |  3 +++
 .../data_loader => data}/tabular_loader.ts      |  4 ++--
 .../data_loader => data}/text_loader.ts         |  4 ++--
 .../src/dataset/data_loader/index.ts            |  3 ---
 discojs/discojs-web/src/imports.ts              |  2 --
 discojs/discojs-web/src/index.ts                |  5 ++---
 discojs/discojs-web/src/memory/memory.ts        |  3 ++-
 server/src/config.ts                            |  2 +-
 server/src/get_server.ts                        |  2 +-
 server/src/router/decentralized/server.ts       |  4 ++--
 server/src/router/federated/server.ts           | 16 +++++++++-------
 server/src/router/server.ts                     |  2 +-
 server/src/router/tasks.ts                      |  4 ++--
 server/src/tasks.ts                             |  4 ++--
 server/tests/client/decentralized.spec.ts       |  3 ++-
 server/tests/client/federated.spec.ts           |  3 ++-
 server/tests/e2e/decentralized.spec.ts          |  5 +++--
 server/tests/e2e/federated.spec.ts              |  9 +++++----
 server/tests/utils.ts                           |  2 +-
 web-client/cypress/e2e/tasks.cy.ts              |  2 +-
 web-client/src/clients.ts                       |  2 +-
 web-client/src/components/data/Data.vue         |  2 +-
 .../src/components/data/DataHarmonization.vue   |  2 +-
 .../data/dataset_input/DatasetInput.vue         |  2 +-
 web-client/src/components/pages/TaskList.vue    |  2 +-
 .../src/components/sidebar/ModelLibrary.vue     |  5 +++--
 .../components/task_creation_form/TaskForm.vue  |  2 +-
 web-client/src/components/testing/Tester.vue    |  5 +++--
 web-client/src/components/testing/Testing.vue   |  9 +++++----
 .../src/components/training/Description.vue     |  2 +-
 web-client/src/components/training/Finished.vue |  6 ++++--
 .../src/components/training/ModelCaching.vue    |  5 +++--
 web-client/src/components/training/Trainer.vue  |  5 +++--
 web-client/src/components/training/Training.vue |  9 +++++----
 .../components/training/TrainingInformation.vue |  2 +-
 web-client/src/store/memory.ts                  |  2 +-
 web-client/src/store/tasks.ts                   |  2 +-
 web-client/src/store/training.ts                |  2 +-
 web-client/src/store/validation.ts              |  2 +-
 56 files changed, 119 insertions(+), 106 deletions(-)
 rename discojs/discojs-node/src/{dataset/data_loader => data}/image_loader.spec.ts (92%)
 rename discojs/discojs-node/src/{dataset/data_loader => data}/image_loader.ts (81%)
 create mode 100644 discojs/discojs-node/src/data/index.ts
 rename discojs/discojs-node/src/{dataset/data_loader => data}/tabular_loader.spec.ts (89%)
 rename discojs/discojs-node/src/{dataset/data_loader => data}/tabular_loader.ts (83%)
 rename discojs/discojs-node/src/{dataset/data_loader => data}/text_loader.ts (100%)
 delete mode 100644 discojs/discojs-node/src/dataset/data_loader/index.ts
 delete mode 100644 discojs/discojs-node/src/imports.ts
 rename discojs/discojs-web/src/{dataset/data_loader => data}/image_loader.ts (62%)
 create mode 100644 discojs/discojs-web/src/data/index.ts
 rename discojs/discojs-web/src/{dataset/data_loader => data}/tabular_loader.ts (67%)
 rename discojs/discojs-web/src/{dataset/data_loader => data}/text_loader.ts (77%)
 delete mode 100644 discojs/discojs-web/src/dataset/data_loader/index.ts
 delete mode 100644 discojs/discojs-web/src/imports.ts

diff --git a/cli/src/args.ts b/cli/src/args.ts
index dd8890089..e8c5b6937 100644
--- a/cli/src/args.ts
+++ b/cli/src/args.ts
@@ -1,6 +1,7 @@
 import { parse } from 'ts-command-line-args'
 import { Map } from 'immutable'
-import { defaultTasks, type Task } from '@epfml/discojs-node'
+import type { Task } from '@epfml/discojs-core'
+import { defaultTasks } from '@epfml/discojs-core'
 
 interface BenchmarkArguments {
   task: Task
diff --git a/cli/src/cli.ts b/cli/src/cli.ts
index d4fc3b248..f483933bf 100644
--- a/cli/src/cli.ts
+++ b/cli/src/cli.ts
@@ -1,6 +1,7 @@
 import { Range } from 'immutable'
 
-import { Disco, TrainingSchemes, type TrainerLog, type data, type Task } from '@epfml/discojs-node'
+import type { TrainerLog, data, Task } from '@epfml/discojs-core'
+import { Disco, TrainingSchemes } from '@epfml/discojs-core'
 
 import { startServer, saveLog } from './utils'
 import { getTaskData } from './data'
diff --git a/cli/src/data.ts b/cli/src/data.ts
index c099ba9bf..c850a7019 100644
--- a/cli/src/data.ts
+++ b/cli/src/data.ts
@@ -3,8 +3,8 @@ import fs from 'node:fs'
 import fs_promises from 'fs/promises'
 import path from 'node:path'
 
-import type { Task, data } from '@epfml/discojs-node'
-import { node } from '@epfml/discojs-node'
+import type { Task, data } from '@epfml/discojs-core'
+import { NodeImageLoader, NodeTabularLoader } from '@epfml/discojs-node'
 
 function filesFromFolder (dir: string, folder: string, fractionToKeep: number): string[] {
   const f = fs.readdirSync(dir + folder)
@@ -35,7 +35,7 @@ async function simplefaceData (task: Task): Promise<data.DataSplit> {
   const labels = filesPerFolder.flatMap((files, index) => Array(files.length).fill(index))
   const files = filesPerFolder.flat()
 
-  return await new node.data.NodeImageLoader(task).loadAll(files, { labels })
+  return await new NodeImageLoader(task).loadAll(files, { labels })
 }
 
 async function cifar10Data (cifar10: Task): Promise<data.DataSplit> {
@@ -43,13 +43,13 @@ async function cifar10Data (cifar10: Task): Promise<data.DataSplit> {
   const files = (await fs_promises.readdir(dir)).map((file) => path.join(dir, file))
   const labels = Range(0, 24).map((label) => (label % 10).toString()).toArray()
 
-  return await new node.data.NodeImageLoader(cifar10).loadAll(files, { labels })
+  return await new NodeImageLoader(cifar10).loadAll(files, { labels })
 }
 
 async function titanicData (titanic: Task): Promise<data.DataSplit> {
   const dir = '../example_training_data/titanic_train.csv'
 
-  const data = await (new node.data.NodeTabularLoader(titanic, ',').loadAll(
+  const data = await (new NodeTabularLoader(titanic, ',').loadAll(
     ['file://'.concat(dir)],
     {
       features: titanic.trainingInformation?.inputColumns,
diff --git a/cli/src/utils.ts b/cli/src/utils.ts
index 42c27b879..3c1f0dc9c 100644
--- a/cli/src/utils.ts
+++ b/cli/src/utils.ts
@@ -1,7 +1,7 @@
 import type http from 'node:http'
 import fs from 'node:fs'
 
-import { type TrainerLog } from '@epfml/discojs-node'
+import type { TrainerLog } from '@epfml/discojs-core'
 import { Disco } from '@epfml/disco-server'
 
 export async function startServer (): Promise<[http.Server, URL]> {
diff --git a/discojs/discojs-core/src/aggregator/secure.spec.ts b/discojs/discojs-core/src/aggregator/secure.spec.ts
index 088a70fb6..0f5f33797 100644
--- a/discojs/discojs-core/src/aggregator/secure.spec.ts
+++ b/discojs/discojs-core/src/aggregator/secure.spec.ts
@@ -1,7 +1,7 @@
 import { List, Set, Range } from 'immutable'
 import { assert } from 'chai'
 
-import { aggregator as aggregators, aggregation, WeightsContainer, defaultTasks } from '@epfml/discojs-node'
+import { aggregator as aggregators, aggregation, WeightsContainer, defaultTasks } from '@epfml/discojs-core'
 
 describe('secret shares test', function () {
   const epsilon = 1e-4
diff --git a/discojs/discojs-core/src/validation/validator.spec.ts b/discojs/discojs-core/src/validation/validator.spec.ts
index 63534cab4..24d3a7db5 100644
--- a/discojs/discojs-core/src/validation/validator.spec.ts
+++ b/discojs/discojs-core/src/validation/validator.spec.ts
@@ -1,10 +1,9 @@
 import { assert } from 'chai'
 import fs from 'fs'
 
-import {
-  type Task, node, Validator, ConsoleLogger, EmptyMemory,
-  client as clients, type data, aggregator, defaultTasks
-} from '@epfml/discojs-node'
+import type { Task, data } from '../..'
+import { Validator, ConsoleLogger, EmptyMemory, client as clients, aggregator, defaultTasks } from '../..'
+import { NodeImageLoader, NodeTabularLoader } from '@epfml/discojs-node'
 
 const simplefaceMock = {
   taskID: 'simple_face',
@@ -32,7 +31,7 @@ describe('validator', () => {
         .map((file: string) => dir + subdir + file))
     const labels = files.flatMap((files, index) => Array(files.length).fill(index))
 
-    const data = (await new node.data.NodeImageLoader(simplefaceMock)
+    const data = (await new NodeImageLoader(simplefaceMock)
       .loadAll(files.flat(), { labels })).train
     const meanAggregator = new aggregator.MeanAggregator(simplefaceMock)
     const client = new clients.Local(new URL('http://localhost:8080'), simplefaceMock, meanAggregator)
@@ -63,7 +62,7 @@ describe('validator', () => {
   it('titanic validator', async () => {
     const titanicTask = defaultTasks.titanic.getTask()
     const files = ['../../example_training_data/titanic_train.csv']
-    const data: data.Data = (await new node.data.NodeTabularLoader(titanicTask, ',').loadAll(files, {
+    const data: data.Data = (await new NodeTabularLoader(titanicTask, ',').loadAll(files, {
       features: titanicTask.trainingInformation.inputColumns,
       labels: titanicTask.trainingInformation.outputColumns,
       shuffle: false
diff --git a/discojs/discojs-core/src/weights/aggregation.spec.ts b/discojs/discojs-core/src/weights/aggregation.spec.ts
index 97774bdc0..fff526cbc 100644
--- a/discojs/discojs-core/src/weights/aggregation.spec.ts
+++ b/discojs/discojs-core/src/weights/aggregation.spec.ts
@@ -1,6 +1,6 @@
 import { assert } from 'chai'
 
-import { WeightsContainer, aggregation } from '@epfml/discojs-node'
+import { WeightsContainer, aggregation } from '.'
 
 describe('weights aggregation', () => {
   it('avg of weights with two operands', async () => {
diff --git a/discojs/discojs-node/src/dataset/data_loader/image_loader.spec.ts b/discojs/discojs-node/src/data/image_loader.spec.ts
similarity index 92%
rename from discojs/discojs-node/src/dataset/data_loader/image_loader.spec.ts
rename to discojs/discojs-node/src/data/image_loader.spec.ts
index c38417b56..dda02de96 100644
--- a/discojs/discojs-node/src/dataset/data_loader/image_loader.spec.ts
+++ b/discojs/discojs-node/src/data/image_loader.spec.ts
@@ -1,11 +1,12 @@
 import { assert, expect } from 'chai'
 import { List, Map, Range } from 'immutable'
 import fs from 'fs'
-
 import tf from '@tensorflow/tfjs'
 import { node as tfNode } from '@tensorflow/tfjs-node'
 
-import { node, type Task } from '../..'
+import type { Task } from '@epfml/discojs-core'
+
+import { ImageLoader } from './image_loader'
 
 const readFilesFromDir = (dir: string): string[] =>
   fs.readdirSync(dir).map((file: string) => dir + file)
@@ -35,8 +36,8 @@ const mnistMock: Task = {
 } as unknown as Task
 
 const LOADERS = {
-  CIFAR10: new node.data.NodeImageLoader(cifar10Mock),
-  MNIST: new node.data.NodeImageLoader(mnistMock)
+  CIFAR10: new ImageLoader(cifar10Mock),
+  MNIST: new ImageLoader(mnistMock)
 }
 const FILES = Map(DIRS).map((readFilesFromDir)).toObject()
 
@@ -93,7 +94,7 @@ describe('image loader', () => {
   })
 
   it('loads samples in order', async () => {
-    const loader = new node.data.NodeImageLoader(cifar10Mock)
+    const loader = new ImageLoader(cifar10Mock)
     const dataset = await ((await loader.loadAll(FILES.CIFAR10, { shuffle: false })).train.dataset).toArray()
 
     List(dataset).zip(List(FILES.CIFAR10))
@@ -105,7 +106,7 @@ describe('image loader', () => {
   })
 
   it('shuffles list', async () => {
-    const loader = new node.data.NodeImageLoader(cifar10Mock)
+    const loader = new ImageLoader(cifar10Mock)
     const list = Range(0, 100_000).toArray()
     const shuffled = [...list]
 
@@ -117,7 +118,7 @@ describe('image loader', () => {
   })
 
   it('shuffles samples', async () => {
-    const loader = new node.data.NodeImageLoader(cifar10Mock)
+    const loader = new ImageLoader(cifar10Mock)
     const dataset = await (await loader.loadAll(FILES.CIFAR10, { shuffle: false })).train.dataset.toArray()
     const shuffled = await (await loader.loadAll(FILES.CIFAR10, { shuffle: true })).train.dataset.toArray()
 
@@ -129,7 +130,7 @@ describe('image loader', () => {
   it('validation split', async () => {
     const validationSplit = 0.2
     const imagesContent = FILES.CIFAR10.map((file) => tfNode.decodeImage(fs.readFileSync(file)))
-    const datasetContent = await new node.data.NodeImageLoader(cifar10Mock)
+    const datasetContent = await new ImageLoader(cifar10Mock)
       .loadAll(FILES.CIFAR10, { shuffle: false, validationSplit })
 
     const trainSize = Math.floor(imagesContent.length * (1 - validationSplit))
diff --git a/discojs/discojs-node/src/dataset/data_loader/image_loader.ts b/discojs/discojs-node/src/data/image_loader.ts
similarity index 81%
rename from discojs/discojs-node/src/dataset/data_loader/image_loader.ts
rename to discojs/discojs-node/src/data/image_loader.ts
index 58b1d7954..2c7fee5db 100644
--- a/discojs/discojs-node/src/dataset/data_loader/image_loader.ts
+++ b/discojs/discojs-node/src/data/image_loader.ts
@@ -4,7 +4,7 @@ import { node } from '@tensorflow/tfjs-node'
 
 import { data } from '@epfml/discojs-core'
 
-export class NodeImageLoader extends data.ImageLoader<string> {
+export class ImageLoader extends data.ImageLoader<string> {
   async readImageFrom (source: string): Promise<tf.Tensor3D> {
     return node.decodeImage(fs.readFileSync(source)) as tf.Tensor3D
   }
diff --git a/discojs/discojs-node/src/data/index.ts b/discojs/discojs-node/src/data/index.ts
new file mode 100644
index 000000000..612a1c891
--- /dev/null
+++ b/discojs/discojs-node/src/data/index.ts
@@ -0,0 +1,2 @@
+export { ImageLoader as NodeImageLoader } from './image_loader'
+export { TabularLoader as NodeTabularLoader } from './tabular_loader'
diff --git a/discojs/discojs-node/src/dataset/data_loader/tabular_loader.spec.ts b/discojs/discojs-node/src/data/tabular_loader.spec.ts
similarity index 89%
rename from discojs/discojs-node/src/dataset/data_loader/tabular_loader.spec.ts
rename to discojs/discojs-node/src/data/tabular_loader.spec.ts
index d1745556e..22c969d9c 100644
--- a/discojs/discojs-node/src/dataset/data_loader/tabular_loader.spec.ts
+++ b/discojs/discojs-node/src/data/tabular_loader.spec.ts
@@ -2,7 +2,9 @@ import { assert, expect } from 'chai'
 import { List } from 'immutable'
 import tf from '@tensorflow/tfjs'
 
-import { node, type Task } from '../..'
+import type { Task } from '@epfml/discojs-core'
+
+import { TabularLoader } from './tabular_loader'
 
 const inputFiles = ['../../example_training_data/titanic_train.csv']
 
@@ -26,7 +28,7 @@ const titanicMock: Task = {
 
 describe('tabular loader', () => {
   it('loads a single sample', async () => {
-    const loaded = new node.data.NodeTabularLoader(titanicMock, ',').loadAll(
+    const loaded = new TabularLoader(titanicMock, ',').loadAll(
       inputFiles,
       {
         features: titanicMock.trainingInformation?.inputColumns,
@@ -51,7 +53,7 @@ describe('tabular loader', () => {
 
   it('shuffles samples', async () => {
     const titanic = titanicMock
-    const loader = new node.data.NodeTabularLoader(titanic, ',')
+    const loader = new TabularLoader(titanic, ',')
     const config = {
       features: titanic.trainingInformation?.inputColumns,
       labels: titanic.trainingInformation?.outputColumns,
diff --git a/discojs/discojs-node/src/dataset/data_loader/tabular_loader.ts b/discojs/discojs-node/src/data/tabular_loader.ts
similarity index 83%
rename from discojs/discojs-node/src/dataset/data_loader/tabular_loader.ts
rename to discojs/discojs-node/src/data/tabular_loader.ts
index 79aec87a2..b2dca6e86 100644
--- a/discojs/discojs-node/src/dataset/data_loader/tabular_loader.ts
+++ b/discojs/discojs-node/src/data/tabular_loader.ts
@@ -2,7 +2,7 @@ import { data as tfData } from '@tensorflow/tfjs-node'
 
 import { data } from '@epfml/discojs-core'
 
-export class NodeTabularLoader extends data.TabularLoader<string> {
+export class TabularLoader extends data.TabularLoader<string> {
   async loadDatasetFrom (source: string, csvConfig: Record<string, unknown>): Promise<data.Dataset> {
     const prefix = 'file://'
     if (source.slice(0, 7) !== prefix) {
diff --git a/discojs/discojs-node/src/dataset/data_loader/text_loader.ts b/discojs/discojs-node/src/data/text_loader.ts
similarity index 100%
rename from discojs/discojs-node/src/dataset/data_loader/text_loader.ts
rename to discojs/discojs-node/src/data/text_loader.ts
diff --git a/discojs/discojs-node/src/dataset/data_loader/index.ts b/discojs/discojs-node/src/dataset/data_loader/index.ts
deleted file mode 100644
index 0ea714444..000000000
--- a/discojs/discojs-node/src/dataset/data_loader/index.ts
+++ /dev/null
@@ -1,2 +0,0 @@
-export { NodeImageLoader } from './image_loader'
-export { NodeTabularLoader } from './tabular_loader'
diff --git a/discojs/discojs-node/src/imports.ts b/discojs/discojs-node/src/imports.ts
deleted file mode 100644
index f4628360c..000000000
--- a/discojs/discojs-node/src/imports.ts
+++ /dev/null
@@ -1 +0,0 @@
-export * as data from './dataset/data_loader'
diff --git a/discojs/discojs-node/src/index.ts b/discojs/discojs-node/src/index.ts
index 8b4108de4..7f165348f 100644
--- a/discojs/discojs-node/src/index.ts
+++ b/discojs/discojs-node/src/index.ts
@@ -1,3 +1 @@
-export * from '@epfml/discojs-core'
-
-export * as node from './imports'
+export * from './data'
diff --git a/discojs/discojs-web/src/dataset/data_loader/image_loader.ts b/discojs/discojs-web/src/data/image_loader.ts
similarity index 62%
rename from discojs/discojs-web/src/dataset/data_loader/image_loader.ts
rename to discojs/discojs-web/src/data/image_loader.ts
index 6cabcc0b7..537dfe26c 100644
--- a/discojs/discojs-web/src/dataset/data_loader/image_loader.ts
+++ b/discojs/discojs-web/src/data/image_loader.ts
@@ -1,8 +1,8 @@
 import tf from '@tensorflow/tfjs'
 
-import { data } from '../..'
+import { data } from '@epfml/discojs-core'
 
-export class WebImageLoader extends data.ImageLoader<File> {
+export class ImageLoader extends data.ImageLoader<File> {
   async readImageFrom (source: File): Promise<tf.Tensor3D> {
     return tf.browser.fromPixels(await createImageBitmap(source))
   }
diff --git a/discojs/discojs-web/src/data/index.ts b/discojs/discojs-web/src/data/index.ts
new file mode 100644
index 000000000..705a83eb9
--- /dev/null
+++ b/discojs/discojs-web/src/data/index.ts
@@ -0,0 +1,3 @@
+export { ImageLoader as WebImageLoader } from './image_loader'
+export { TabularLoader as WebTabularLoader } from './tabular_loader'
+export { TextLoader as WebTextLoader } from './text_loader'
diff --git a/discojs/discojs-web/src/dataset/data_loader/tabular_loader.ts b/discojs/discojs-web/src/data/tabular_loader.ts
similarity index 67%
rename from discojs/discojs-web/src/dataset/data_loader/tabular_loader.ts
rename to discojs/discojs-web/src/data/tabular_loader.ts
index f4ab7425c..edb9e7d6c 100644
--- a/discojs/discojs-web/src/dataset/data_loader/tabular_loader.ts
+++ b/discojs/discojs-web/src/data/tabular_loader.ts
@@ -1,8 +1,8 @@
 import tf from '@tensorflow/tfjs'
 
-import { data } from '../..'
+import { data } from '@epfml/discojs-core'
 
-export class WebTabularLoader extends data.TabularLoader<File> {
+export class TabularLoader extends data.TabularLoader<File> {
   async loadDatasetFrom (source: File, csvConfig: Record<string, unknown>): Promise<data.Dataset> {
     return new tf.data.CSVDataset(new tf.data.FileDataSource(source), csvConfig)
   }
diff --git a/discojs/discojs-web/src/dataset/data_loader/text_loader.ts b/discojs/discojs-web/src/data/text_loader.ts
similarity index 77%
rename from discojs/discojs-web/src/dataset/data_loader/text_loader.ts
rename to discojs/discojs-web/src/data/text_loader.ts
index eab2c98b5..702368ca3 100644
--- a/discojs/discojs-web/src/dataset/data_loader/text_loader.ts
+++ b/discojs/discojs-web/src/data/text_loader.ts
@@ -1,8 +1,8 @@
 import tf from '@tensorflow/tfjs'
 
-import { data } from '../..'
+import { data } from '@epfml/discojs-core'
 
-export class WebTextLoader extends data.TextLoader<File> {
+export class TextLoader extends data.TextLoader<File> {
   async loadDatasetFrom (source: File, config?: Record<string, unknown>): Promise<data.Dataset> {
     const file = new tf.data.FileDataSource(source)
     if (config !== undefined) {
diff --git a/discojs/discojs-web/src/dataset/data_loader/index.ts b/discojs/discojs-web/src/dataset/data_loader/index.ts
deleted file mode 100644
index 2dd904507..000000000
--- a/discojs/discojs-web/src/dataset/data_loader/index.ts
+++ /dev/null
@@ -1,3 +0,0 @@
-export { WebImageLoader } from './image_loader'
-export { WebTabularLoader } from './tabular_loader'
-export { WebTextLoader } from './text_loader'
diff --git a/discojs/discojs-web/src/imports.ts b/discojs/discojs-web/src/imports.ts
deleted file mode 100644
index 8b334abde..000000000
--- a/discojs/discojs-web/src/imports.ts
+++ /dev/null
@@ -1,2 +0,0 @@
-export * as data from './dataset/data_loader'
-export { IndexedDB } from './memory'
diff --git a/discojs/discojs-web/src/index.ts b/discojs/discojs-web/src/index.ts
index aef4db385..f55da3c7a 100644
--- a/discojs/discojs-web/src/index.ts
+++ b/discojs/discojs-web/src/index.ts
@@ -1,3 +1,2 @@
-export * from '@epfml/discojs-core'
-
-export * as browser from './imports'
+export * from './data'
+export * from './memory'
diff --git a/discojs/discojs-web/src/memory/memory.ts b/discojs/discojs-web/src/memory/memory.ts
index d821812c7..ab7296d20 100644
--- a/discojs/discojs-web/src/memory/memory.ts
+++ b/discojs/discojs-web/src/memory/memory.ts
@@ -10,7 +10,8 @@ import { Map } from 'immutable'
 import path from 'path'
 import * as tf from '@tensorflow/tfjs'
 
-import { Memory, ModelType, type Path, type ModelInfo, type ModelSource } from '..'
+import type { Path, ModelInfo, ModelSource } from '@epfml/discojs-core'
+import { Memory, ModelType } from '@epfml/discojs-core'
 
 export class IndexedDB extends Memory {
   pathFor (source: ModelSource): Path {
diff --git a/server/src/config.ts b/server/src/config.ts
index 9c295acc0..b1e27fcd2 100644
--- a/server/src/config.ts
+++ b/server/src/config.ts
@@ -1,6 +1,6 @@
 import path from 'path'
 
-import { type Path, type TaskID } from '@epfml/discojs-node'
+import type { Path, TaskID } from '@epfml/discojs-core'
 
 export class Config {
   public readonly serverUrl: URL
diff --git a/server/src/get_server.ts b/server/src/get_server.ts
index a124f77f8..68bf491ef 100644
--- a/server/src/get_server.ts
+++ b/server/src/get_server.ts
@@ -3,7 +3,7 @@ import express from 'express'
 import expressWS from 'express-ws'
 import type tf from '@tensorflow/tfjs'
 
-import type { Task, TaskProvider } from '@epfml/discojs-node'
+import type { Task, TaskProvider } from '@epfml/discojs-core'
 
 import { CONFIG } from './config'
 import { Router } from './router'
diff --git a/server/src/router/decentralized/server.ts b/server/src/router/decentralized/server.ts
index d779720a9..99f780a5f 100644
--- a/server/src/router/decentralized/server.ts
+++ b/server/src/router/decentralized/server.ts
@@ -7,8 +7,8 @@ import type { ParsedQs } from 'qs'
 import { Map, Set } from 'immutable'
 import type tf from '@tensorflow/tfjs'
 
-import type { Task, TaskID } from '@epfml/discojs-node'
-import { client } from '@epfml/discojs-node'
+import type { Task, TaskID } from '@epfml/discojs-core'
+import { client } from '@epfml/discojs-core'
 
 import { Server } from '../server'
 
diff --git a/server/src/router/federated/server.ts b/server/src/router/federated/server.ts
index 5444a15b1..fc3c5cf96 100644
--- a/server/src/router/federated/server.ts
+++ b/server/src/router/federated/server.ts
@@ -5,17 +5,19 @@ import { List, Map } from 'immutable'
 import msgpack from 'msgpack-lite'
 import type tf from '@tensorflow/tfjs'
 
+import type {
+  Task,
+  TaskID,
+  WeightsContainer,
+  MetadataKey,
+  MetadataValue
+} from '@epfml/discojs-core'
 import {
   client,
   serialization,
   AsyncInformant,
-  type Task,
-  type TaskID,
-  aggregator as aggregators,
-  type WeightsContainer,
-  type MetadataKey,
-  type MetadataValue
-} from '@epfml/discojs-node'
+  aggregator as aggregators
+} from '@epfml/discojs-core'
 
 import { Server } from '../server'
 
diff --git a/server/src/router/server.ts b/server/src/router/server.ts
index 7ee7ba9fb..d2a6fe023 100644
--- a/server/src/router/server.ts
+++ b/server/src/router/server.ts
@@ -3,7 +3,7 @@ import type expressWS from 'express-ws'
 import type WebSocket from 'ws'
 import type tf from '@tensorflow/tfjs'
 
-import type { Task } from '@epfml/discojs-node'
+import type { Task } from '@epfml/discojs-core'
 
 import type { TasksAndModels } from '../tasks'
 
diff --git a/server/src/router/tasks.ts b/server/src/router/tasks.ts
index 51bc58fe3..2eef48bd8 100644
--- a/server/src/router/tasks.ts
+++ b/server/src/router/tasks.ts
@@ -3,8 +3,8 @@ import express from 'express'
 import { Set } from 'immutable'
 import type tf from '@tensorflow/tfjs'
 
-import type { Task, TaskID } from '@epfml/discojs-node'
-import { serialization, isTask } from '@epfml/discojs-node'
+import type { Task, TaskID } from '@epfml/discojs-core'
+import { serialization, isTask } from '@epfml/discojs-core'
 
 import type { Config } from '../config'
 import type { TasksAndModels } from '../tasks'
diff --git a/server/src/tasks.ts b/server/src/tasks.ts
index 25f23b7f0..e141fb333 100644
--- a/server/src/tasks.ts
+++ b/server/src/tasks.ts
@@ -4,8 +4,8 @@ import fs from 'node:fs'
 import tf from '@tensorflow/tfjs'
 import '@tensorflow/tfjs-node'
 
-import type { Task, Path, Digest, TaskProvider } from '@epfml/discojs-node'
-import { isTaskProvider, defaultTasks } from '@epfml/discojs-node'
+import type { Task, Path, Digest, TaskProvider } from '@epfml/discojs-core'
+import { isTaskProvider, defaultTasks } from '@epfml/discojs-core'
 
 // default tasks and added ones
 // register 'taskAndModel' event to get tasks
diff --git a/server/tests/client/decentralized.spec.ts b/server/tests/client/decentralized.spec.ts
index fc74aafb2..48bf560d0 100644
--- a/server/tests/client/decentralized.spec.ts
+++ b/server/tests/client/decentralized.spec.ts
@@ -1,6 +1,7 @@
 import type * as http from 'http'
 
-import { aggregator as aggregators, client as clients, type Task, defaultTasks } from '@epfml/discojs-node'
+import type { Task } from '@epfml/discojs-core'
+import { aggregator as aggregators, client as clients, defaultTasks } from '@epfml/discojs-core'
 
 import { getClient, startServer } from '../utils'
 
diff --git a/server/tests/client/federated.spec.ts b/server/tests/client/federated.spec.ts
index d6b4f3387..c57772c72 100644
--- a/server/tests/client/federated.spec.ts
+++ b/server/tests/client/federated.spec.ts
@@ -1,6 +1,7 @@
 import type * as http from 'http'
 
-import { aggregator as aggregators, client as clients, informant, defaultTasks, type Task } from '@epfml/discojs-node'
+import type { Task } from '@epfml/discojs-core'
+import { aggregator as aggregators, client as clients, informant, defaultTasks } from '@epfml/discojs-core'
 
 import { getClient, startServer } from '../utils'
 
diff --git a/server/tests/e2e/decentralized.spec.ts b/server/tests/e2e/decentralized.spec.ts
index 769415807..67ad64a40 100644
--- a/server/tests/e2e/decentralized.spec.ts
+++ b/server/tests/e2e/decentralized.spec.ts
@@ -2,9 +2,10 @@ import type { Server } from 'node:http'
 import { List } from 'immutable'
 import { assert } from 'chai'
 
+import type { Task } from '@epfml/discojs-core'
 import {
-  aggregator as aggregators, informant as informants, client as clients, type Task, WeightsContainer, defaultTasks, aggregation
-} from '@epfml/discojs-node'
+  aggregator as aggregators, informant as informants, client as clients, WeightsContainer, defaultTasks, aggregation
+} from '@epfml/discojs-core'
 
 import { getClient, startServer } from '../utils'
 
diff --git a/server/tests/e2e/federated.spec.ts b/server/tests/e2e/federated.spec.ts
index 53f05359f..095baf2e7 100644
--- a/server/tests/e2e/federated.spec.ts
+++ b/server/tests/e2e/federated.spec.ts
@@ -5,9 +5,10 @@ import { Range } from 'immutable'
 import { assert } from 'chai'
 
 import {
-  WeightsContainer, node, Disco, TrainingSchemes, client as clients,
+  WeightsContainer, Disco, TrainingSchemes, client as clients,
   aggregator as aggregators, informant, defaultTasks
-} from '@epfml/discojs-node'
+} from '@epfml/discojs-core'
+import { NodeImageLoader, NodeTabularLoader } from '@epfml/discojs-node'
 
 import { getClient, startServer } from '../utils'
 
@@ -31,7 +32,7 @@ describe('end-to-end federated', function () {
 
     const cifar10Task = defaultTasks.cifar10.getTask()
 
-    const data = await new node.data.NodeImageLoader(cifar10Task).loadAll(files, { labels })
+    const data = await new NodeImageLoader(cifar10Task).loadAll(files, { labels })
 
     const aggregator = new aggregators.MeanAggregator(cifar10Task)
     const client = await getClient(clients.federated.FederatedClient, server, cifar10Task, aggregator)
@@ -51,7 +52,7 @@ describe('end-to-end federated', function () {
 
     const titanicTask = defaultTasks.titanic.getTask()
     titanicTask.trainingInformation.epochs = 5
-    const data = await (new node.data.NodeTabularLoader(titanicTask, ',').loadAll(
+    const data = await (new NodeTabularLoader(titanicTask, ',').loadAll(
       files,
       {
         features: titanicTask.trainingInformation.inputColumns,
diff --git a/server/tests/utils.ts b/server/tests/utils.ts
index aceca4432..de71f9a91 100644
--- a/server/tests/utils.ts
+++ b/server/tests/utils.ts
@@ -1,6 +1,6 @@
 import type { Server } from 'node:http'
 
-import type { aggregator, client, Task } from '@epfml/discojs-node'
+import type { aggregator, client, Task } from '@epfml/discojs-core'
 
 import { runDefaultServer } from '../src/get_server'
 
diff --git a/web-client/cypress/e2e/tasks.cy.ts b/web-client/cypress/e2e/tasks.cy.ts
index 37d961342..b1b8030a3 100644
--- a/web-client/cypress/e2e/tasks.cy.ts
+++ b/web-client/cypress/e2e/tasks.cy.ts
@@ -1,5 +1,5 @@
 /* eslint-disable no-undef */
-import { defaultTasks } from '@epfml/discojs'
+import { defaultTasks } from '@epfml/discojs-core'
 
 // most basic disco tasks
 export const TASK_LIST = [
diff --git a/web-client/src/clients.ts b/web-client/src/clients.ts
index 2eae7078d..6e890cc78 100644
--- a/web-client/src/clients.ts
+++ b/web-client/src/clients.ts
@@ -1,4 +1,4 @@
-import { client as clients, aggregator as aggregators, Task, TrainingSchemes } from '@epfml/discojs'
+import { client as clients, aggregator as aggregators, Task, TrainingSchemes } from '@epfml/discojs-core'
 
 import { CONFIG } from './config'
 
diff --git a/web-client/src/components/data/Data.vue b/web-client/src/components/data/Data.vue
index a910d53ca..ba4e0eff2 100644
--- a/web-client/src/components/data/Data.vue
+++ b/web-client/src/components/data/Data.vue
@@ -12,7 +12,7 @@
 <script setup lang="ts">
 import { defineProps } from 'vue'
 
-import { data, Task } from '@epfml/discojs'
+import { data, Task } from '@epfml/discojs-core'
 
 import DatasetInput from '@/components/data/dataset_input/DatasetInput.vue'
 import DataHarmonization from '@/components/data/DataHarmonization.vue'
diff --git a/web-client/src/components/data/DataHarmonization.vue b/web-client/src/components/data/DataHarmonization.vue
index c4bd5fad3..1f6ce64f5 100644
--- a/web-client/src/components/data/DataHarmonization.vue
+++ b/web-client/src/components/data/DataHarmonization.vue
@@ -1,7 +1,7 @@
 <script setup lang="ts">
 import { defineProps } from 'vue'
 
-import { Task } from '@epfml/discojs'
+import { Task } from '@epfml/discojs-core'
 
 import DropdownCard from '@/components/containers/DropdownCard.vue'
 
diff --git a/web-client/src/components/data/dataset_input/DatasetInput.vue b/web-client/src/components/data/dataset_input/DatasetInput.vue
index 6a797ef42..c50c8a50f 100644
--- a/web-client/src/components/data/dataset_input/DatasetInput.vue
+++ b/web-client/src/components/data/dataset_input/DatasetInput.vue
@@ -144,7 +144,7 @@
 import { computed, defineProps, ref } from 'vue'
 import * as d3 from 'd3'
 
-import { data, Task } from '@epfml/discojs'
+import { data, Task } from '@epfml/discojs-core'
 
 import Upload from '@/assets/svg/Upload.vue'
 import IconCard from '@/components/containers/IconCard.vue'
diff --git a/web-client/src/components/pages/TaskList.vue b/web-client/src/components/pages/TaskList.vue
index cd4fa53e9..5854a6f3c 100644
--- a/web-client/src/components/pages/TaskList.vue
+++ b/web-client/src/components/pages/TaskList.vue
@@ -122,7 +122,7 @@ import { computed, reactive, ref } from 'vue'
 import { useRouter } from 'vue-router'
 import { storeToRefs } from 'pinia'
 
-import { Task } from '@epfml/discojs'
+import { Task } from '@epfml/discojs-core'
 
 import { useTasksStore } from '@/store/tasks'
 import { useTrainingStore } from '@/store/training'
diff --git a/web-client/src/components/sidebar/ModelLibrary.vue b/web-client/src/components/sidebar/ModelLibrary.vue
index 24394420b..44ca06027 100644
--- a/web-client/src/components/sidebar/ModelLibrary.vue
+++ b/web-client/src/components/sidebar/ModelLibrary.vue
@@ -93,7 +93,8 @@
 import { defineComponent } from 'vue'
 import { mapStores } from 'pinia'
 
-import { browser, Memory, EmptyMemory, Path, ModelType } from '@epfml/discojs'
+import { Memory, EmptyMemory, Path, ModelType } from '@epfml/discojs-core'
+import { IndexedDB } from '@epfml/discojs'
 
 import { useMemoryStore } from '@/store/memory'
 import { useValidationStore } from '@/store/validation'
@@ -121,7 +122,7 @@ export default defineComponent({
     ...mapStores(useMemoryStore, useValidationStore),
 
     memory (): Memory {
-      return this.memoryStore.useIndexedDB ? new browser.IndexedDB() : new EmptyMemory()
+      return this.memoryStore.useIndexedDB ? new IndexedDB() : new EmptyMemory()
     }
   },
   async mounted (): Promise<void> {
diff --git a/web-client/src/components/task_creation_form/TaskForm.vue b/web-client/src/components/task_creation_form/TaskForm.vue
index be92194d8..69bc5c09e 100644
--- a/web-client/src/components/task_creation_form/TaskForm.vue
+++ b/web-client/src/components/task_creation_form/TaskForm.vue
@@ -158,7 +158,7 @@ import { Form as VeeForm, ErrorMessage } from 'vee-validate'
 import { List, Map } from 'immutable'
 import * as tf from '@tensorflow/tfjs'
 
-import { Task, pushTask } from '@epfml/discojs'
+import { Task, pushTask } from '@epfml/discojs-core'
 
 import { sections, modelCompileData, privacyParameters, FormField, FormSection } from '@/task_creation_form'
 import { useToaster } from '@/composables/toaster'
diff --git a/web-client/src/components/testing/Tester.vue b/web-client/src/components/testing/Tester.vue
index 1dc57b3e2..9bf1941bf 100644
--- a/web-client/src/components/testing/Tester.vue
+++ b/web-client/src/components/testing/Tester.vue
@@ -237,7 +237,8 @@
 import { computed, defineProps, ref } from 'vue'
 import { storeToRefs } from 'pinia'
 
-import { browser, data, ConsoleLogger, EmptyMemory, Memory, Task, Validator, LabelTypeEnum } from '@epfml/discojs'
+import { data, ConsoleLogger, EmptyMemory, Memory, Task, Validator, LabelTypeEnum } from '@epfml/discojs-core'
+import { IndexedDB } from '@epfml/discojs'
 
 import { useMemoryStore } from '@/store/memory'
 import { useValidationStore } from '@/store/validation'
@@ -287,7 +288,7 @@ const isImageTaskType = computed<boolean>(() =>
 const isPolygonMapVisualization = computed<boolean>(() =>
   props.task.displayInformation.labelDisplay.labelType === LabelTypeEnum.POLYGON_MAP)
 
-const memory = computed<Memory>(() => useIndexedDB ? new browser.IndexedDB() : new EmptyMemory())
+const memory = computed<Memory>(() => useIndexedDB ? new IndexedDB() : new EmptyMemory())
 const accuracyData = computed<number[]>(() => {
   const r = validator.value?.accuracyData
   return r !== undefined ? r.toArray() : [0]
diff --git a/web-client/src/components/testing/Testing.vue b/web-client/src/components/testing/Testing.vue
index 08c01927d..1447ef674 100644
--- a/web-client/src/components/testing/Testing.vue
+++ b/web-client/src/components/testing/Testing.vue
@@ -162,7 +162,8 @@ import { RouterLink } from 'vue-router'
 import { storeToRefs } from 'pinia'
 import { List } from 'immutable'
 
-import { browser, EmptyMemory, Memory, Path, Task, data, ModelType, client as clients } from '@epfml/discojs'
+import { EmptyMemory, Memory, Path, Task, data, ModelType, client as clients } from '@epfml/discojs-core'
+import { IndexedDB, WebImageLoader, WebTabularLoader } from '@epfml/discojs'
 
 import { CONFIG } from '@/config'
 import { useMemoryStore } from '@/store/memory'
@@ -199,7 +200,7 @@ const federatedTasks = computed<List<Task>>(() =>
   tasksStore.tasks.filter((t) => t.trainingInformation.scheme === 'Federated').toList())
 
 const memory = computed<Memory>(() =>
-  memoryStore.useIndexedDB ? new browser.IndexedDB() : new EmptyMemory())
+  memoryStore.useIndexedDB ? new IndexedDB() : new EmptyMemory())
 
 const datasetBuilder = computed<data.DatasetBuilder<File> | undefined>(() => {
   if (currentTask.value === undefined) {
@@ -208,10 +209,10 @@ const datasetBuilder = computed<data.DatasetBuilder<File> | undefined>(() => {
   let dataLoader: data.DataLoader<File>
   switch (currentTask.value.trainingInformation.dataType) {
     case 'tabular':
-      dataLoader = new browser.data.WebTabularLoader(currentTask.value, ',')
+      dataLoader = new WebTabularLoader(currentTask.value, ',')
       break
     case 'image':
-      dataLoader = new browser.data.WebImageLoader(currentTask.value)
+      dataLoader = new WebImageLoader(currentTask.value)
       break
     default:
       throw new Error('not implemented')
diff --git a/web-client/src/components/training/Description.vue b/web-client/src/components/training/Description.vue
index 0ef65d95e..a2f9ae36c 100644
--- a/web-client/src/components/training/Description.vue
+++ b/web-client/src/components/training/Description.vue
@@ -90,7 +90,7 @@
 <script lang="ts" setup>
 import { computed, defineProps } from 'vue'
 
-import { Task } from '@epfml/discojs'
+import { Task } from '@epfml/discojs-core'
 
 import { trainingInformation, privacyParameters, modelCompileData, FormField, FormSection } from '@/task_creation_form'
 import ModelCaching from './ModelCaching.vue'
diff --git a/web-client/src/components/training/Finished.vue b/web-client/src/components/training/Finished.vue
index 6cf99ffc5..cf207a2ef 100644
--- a/web-client/src/components/training/Finished.vue
+++ b/web-client/src/components/training/Finished.vue
@@ -38,7 +38,9 @@
 <script setup lang="ts">
 import { defineProps, computed } from 'vue'
 import { useRouter } from 'vue-router'
-import { browser, EmptyMemory, Memory, ModelType, Task, ModelInfo } from '@epfml/discojs'
+
+import { EmptyMemory, Memory, ModelType, Task, ModelInfo } from '@epfml/discojs-core'
+import { IndexedDB } from '@epfml/discojs'
 
 import { useMemoryStore } from '@/store/memory'
 import { useValidationStore } from '@/store/validation'
@@ -56,7 +58,7 @@ const toast = useToaster() // TODO: use this.$toaster
 
 const props = defineProps<Props>()
 
-const memory = computed<Memory>(() => memoryStore.useIndexedDB ? new browser.IndexedDB() : new EmptyMemory())
+const memory = computed<Memory>(() => memoryStore.useIndexedDB ? new IndexedDB() : new EmptyMemory())
 const modelInfo = computed<ModelInfo>(() => {
   return {
     type: ModelType.WORKING,
diff --git a/web-client/src/components/training/ModelCaching.vue b/web-client/src/components/training/ModelCaching.vue
index 3ce013d6b..9037fb443 100644
--- a/web-client/src/components/training/ModelCaching.vue
+++ b/web-client/src/components/training/ModelCaching.vue
@@ -134,7 +134,8 @@
 <script lang="ts">
 import { mapStores } from 'pinia'
 
-import { browser, EmptyMemory, Memory, ModelType, isTask, TrainingSchemes, ModelInfo } from '@epfml/discojs'
+import { EmptyMemory, Memory, ModelType, isTask, TrainingSchemes, ModelInfo } from '@epfml/discojs-core'
+import { IndexedDB } from '@epfml/discojs'
 
 import { getClient } from '@/clients'
 import Clock from '@/assets/svg/Clock.vue'
@@ -177,7 +178,7 @@ export default {
       )
     },
     memory (): Memory {
-      return this.memoryStore.useIndexedDB ? new browser.IndexedDB() : new EmptyMemory()
+      return this.memoryStore.useIndexedDB ? new IndexedDB() : new EmptyMemory()
     },
     modelInfo (): ModelInfo {
       return {
diff --git a/web-client/src/components/training/Trainer.vue b/web-client/src/components/training/Trainer.vue
index b9040ac21..5c0f12be6 100644
--- a/web-client/src/components/training/Trainer.vue
+++ b/web-client/src/components/training/Trainer.vue
@@ -48,7 +48,8 @@
 import { defineComponent } from 'vue'
 import { mapStores } from 'pinia'
 
-import { browser, data, EmptyMemory, isTask, informant, TrainingInformant, TrainingSchemes, Disco, Memory, client as clients } from '@epfml/discojs'
+import { data, EmptyMemory, isTask, informant, TrainingInformant, TrainingSchemes, Disco, Memory, client as clients } from '@epfml/discojs-core'
+import { IndexedDB } from '@epfml/discojs'
 
 import { useMemoryStore } from '@/store/memory'
 // TODO @s314cy: move to discojs-core/src/client/get.ts
@@ -91,7 +92,7 @@ export default defineComponent({
       return getClient(this.scheme, this.task)
     },
     memory (): Memory {
-      return this.memoryStore.useIndexedDB ? new browser.IndexedDB() : new EmptyMemory()
+      return this.memoryStore.useIndexedDB ? new IndexedDB() : new EmptyMemory()
     },
     disco (): Disco {
       return new Disco(
diff --git a/web-client/src/components/training/Training.vue b/web-client/src/components/training/Training.vue
index 19cde6bc9..cb70711d4 100644
--- a/web-client/src/components/training/Training.vue
+++ b/web-client/src/components/training/Training.vue
@@ -26,7 +26,8 @@
 import { computed, onMounted, ref, defineProps } from 'vue'
 import { useRouter } from 'vue-router'
 
-import { browser, data, TaskID } from '@epfml/discojs'
+import { data, TaskID } from '@epfml/discojs-core'
+import { WebImageLoader, WebTabularLoader, WebTextLoader } from '@epfml/discojs'
 
 import { useTrainingStore } from '@/store/training'
 import { useTasksStore } from '@/store/tasks'
@@ -56,13 +57,13 @@ const datasetBuilder = computed(() => {
   let dataLoader: data.DataLoader<File>
   switch (task.value.trainingInformation.dataType) {
     case 'tabular':
-      dataLoader = new browser.data.WebTabularLoader(task.value)
+      dataLoader = new WebTabularLoader(task.value)
       break
     case 'image':
-      dataLoader = new browser.data.WebImageLoader(task.value)
+      dataLoader = new WebImageLoader(task.value)
       break
     case 'text':
-      dataLoader = new browser.data.WebTextLoader(task.value)
+      dataLoader = new WebTextLoader(task.value)
       break
     default:
       throw new Error('not implemented')
diff --git a/web-client/src/components/training/TrainingInformation.vue b/web-client/src/components/training/TrainingInformation.vue
index 50d406e96..2e2b53969 100644
--- a/web-client/src/components/training/TrainingInformation.vue
+++ b/web-client/src/components/training/TrainingInformation.vue
@@ -134,7 +134,7 @@
 <script setup lang="ts">
 import { defineProps, computed } from 'vue'
 
-import { TrainingInformant } from '@epfml/discojs'
+import { TrainingInformant } from '@epfml/discojs-core'
 
 import { chartOptions } from '@/charts'
 import IconCardSmall from '@/components/containers/IconCardSmall.vue'
diff --git a/web-client/src/store/memory.ts b/web-client/src/store/memory.ts
index 0dd04b7bc..1fe1aa5e6 100644
--- a/web-client/src/store/memory.ts
+++ b/web-client/src/store/memory.ts
@@ -4,7 +4,7 @@ import { ref, shallowRef } from 'vue'
 import { Map } from 'immutable'
 import * as tf from '@tensorflow/tfjs'
 
-import { ModelInfo, ModelType, Path } from '@epfml/discojs'
+import { ModelInfo, ModelType, Path } from '@epfml/discojs-core'
 
 interface ModelMetadata extends ModelInfo {
   date: string
diff --git a/web-client/src/store/tasks.ts b/web-client/src/store/tasks.ts
index aa71da032..4abc4c33f 100644
--- a/web-client/src/store/tasks.ts
+++ b/web-client/src/store/tasks.ts
@@ -2,7 +2,7 @@ import { defineStore } from 'pinia'
 import { shallowRef } from 'vue'
 import { Map } from 'immutable'
 import { useToaster } from '@/composables/toaster'
-import { TaskID, Task, fetchTasks } from '@epfml/discojs'
+import { TaskID, Task, fetchTasks } from '@epfml/discojs-core'
 
 import { CONFIG } from '@/config'
 import { useTrainingStore } from './training'
diff --git a/web-client/src/store/training.ts b/web-client/src/store/training.ts
index 957570c8e..e4bcbf30b 100644
--- a/web-client/src/store/training.ts
+++ b/web-client/src/store/training.ts
@@ -2,7 +2,7 @@ import { defineStore } from 'pinia'
 import { computed, ref, shallowRef } from 'vue'
 import { Map } from 'immutable'
 
-import { TaskID } from '@epfml/discojs'
+import { TaskID } from '@epfml/discojs-core'
 
 export const useTrainingStore = defineStore('training', () => {
   const steps = shallowRef<Map<TaskID, number>>(Map())
diff --git a/web-client/src/store/validation.ts b/web-client/src/store/validation.ts
index 6c58813dc..ee96ae206 100644
--- a/web-client/src/store/validation.ts
+++ b/web-client/src/store/validation.ts
@@ -1,7 +1,7 @@
 import { defineStore } from 'pinia'
 import { ref } from 'vue'
 
-import { Path } from '@epfml/discojs'
+import { Path } from '@epfml/discojs-core'
 
 export const useValidationStore = defineStore('validation', () => {
   const state = ref(false)