From b425596208636c850dfda662f5ea783d4709c54d 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] discojs-core: don't reexport tfjs

---
 cli/src/data.ts                                   |  4 +++-
 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 ++-
 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                               |  4 +++-
 .../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 +-
 65 files changed, 215 insertions(+), 120 deletions(-)

diff --git a/cli/src/data.ts b/cli/src/data.ts
index b19617162..3025fed44 100644
--- a/cli/src/data.ts
+++ b/cli/src/data.ts
@@ -2,8 +2,10 @@ import { Range } from 'immutable'
 import fs from 'node:fs'
 import fs_promises from 'fs/promises'
 import path from 'node:path'
+import tf from '@tensorflow/tfjs-node'
 
-import { tf, node, data, type Task } from '@epfml/discojs-node'
+import type { Task } from '@epfml/discojs-node'
+import { node, data } 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/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..e787a68a8 100644
--- a/server/src/tasks.ts
+++ b/server/src/tasks.ts
@@ -1,8 +1,10 @@
 import { List, Set } from 'immutable'
 import { createHash } from 'node:crypto'
 import fs from 'node:fs'
+import tf from '@tensorflow/tfjs'
 
-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[]