Typescript & tfjs-node: How to work with the node-gpu version alongside plain node?

Question

I have issues running my tfjs node app with and without the GPU library. In plain javascript, as in the examples - they simply require() either @tensorflow/tfjs-node or @tensorflow/tfjs-node-gpu. However, in the Typescript setup I have, require are not allowed. And I need to get the types out of the library. Hence, I "statically" import one of the versions around the place. Then, when the user chooses --gpu or not, i dynamically import() the given module. The problem is that the "static imports" around the place has already caused a module load, and the dynamic load of the other module causes a crash along these lines:

E tensorflow/core/lib/monitoring/collection_registry.cc:81] Cannot register 2 metrics with the same name: /tensorflow/core/op_expansion/node_counter

How can I allow the user to select the right module to use, and still get the typing I need compile time?

A little side question: Can I just always use tfjs-node-gpu? This seems to work on my Mac, where I don't have Cuda, nor anything?

Answer 1

The only way I have found, is to statically import a given set of common types, and use the return type of the dynamic loader modules for the full TF library. Then add this to a wrapper module, like this:

import {Tensor, SymbolicTensor, LayersModel} from "@tensorflow/tfjs-node"

// Exporting these types with the `type` keyword might make a difference in what content is actually imported above. 
export type {Tensor, SymbolicTensor, LayersModel}

export const loadTfjsNode = () => import("@tensorflow/tfjs-node")
export const loadTfjsGpu = () => import("@tensorflow/tfjs-node-gpu")

const loadTfjs = loadTfjsNode
export type ITensorflow = Awaited<ReturnType<typeof loadTfjs>>

This works some some extent, but I can't really access all types this way. I also have to pass the loaded tf module around to all my modules, which is OK in some ways, not OK in others.