@hoff97/tensor-js
Version:
PyTorch like deep learning inferrence library
48 lines • 1.61 kB
JavaScript
export class Dispatcher {
constructor(getOp, minCallsToCompile = 2) {
this.getOp = getOp;
this.minCallsToCompile = minCallsToCompile;
this.opDict = {};
}
getDefault(dtype) {
const str = `default-${dtype}`;
if (this.opDict[str] === undefined) {
const op = this.getOp(dtype);
op.compile({});
this.opDict[str] = {
infoString: str,
numCalls: 0,
operation: op,
};
}
return this.opDict[str];
}
calc(input, dtype) {
const defaultOp = this.getDefault(dtype);
//@ts-ignore
const compileInfoString = defaultOp.operation.getInputInfoString(input);
if (this.opDict[compileInfoString] === undefined) {
this.opDict[compileInfoString] = {
infoString: compileInfoString,
numCalls: 0,
};
}
const opInfo = this.opDict[compileInfoString];
opInfo.numCalls++;
if (opInfo.numCalls >= this.minCallsToCompile) {
if (opInfo.operation === undefined) {
opInfo.operation = this.getOp(dtype);
//@ts-ignore
const compileInfo = defaultOp.operation.getCompilationInfo(input);
opInfo.operation.compile(compileInfo);
}
return opInfo.operation.calc(input);
}
else {
defaultOp.numCalls++;
//@ts-ignore
return defaultOp.operation.calc(input);
}
}
}
//# sourceMappingURL=dispatcher.js.map