@hoff97/tensor-js
Version:
PyTorch like deep learning inferrence library
59 lines • 2.01 kB
JavaScript
export class Optimization {
}
export class SequenceOptimization extends Optimization {
constructor(nodeTypes) {
super();
this.nodeTypes = nodeTypes;
}
findApplications(model) {
const results = [];
const nodes = model.getNodes();
for (const nodeId in Object.keys(nodes)) {
const node = nodes[nodeId];
if (node !== undefined && node.getType() === this.nodeTypes[0]) {
const app = this.checkApplication(model, nodeId);
if (app !== undefined) {
// eslint-disable-next-line @typescript-eslint/no-explicit-any
results.push(app);
}
}
}
return results;
}
checkApplication(model, nodeId) {
const nodes = model.getNodes();
const nodeSeq = [nodeId];
// eslint-disable-next-line @typescript-eslint/no-explicit-any
let lastNode = nodes[nodeId];
const nodeInstances = [lastNode];
for (let i = 1; i < this.nodeTypes.length; i++) {
const nextNodeId = model.getNodeWithInput(lastNode.outputs[0]);
if (nextNodeId !== undefined) {
const nextNode = nodes[nextNodeId];
if (nextNode.getType() === this.nodeTypes[i]) {
// eslint-disable-next-line @typescript-eslint/no-explicit-any
nodeSeq.push(nextNodeId);
lastNode = nextNode;
nodeInstances.push(lastNode);
}
else {
return undefined;
}
}
else {
return undefined;
}
}
if (this.canApply(nodeInstances)) {
return nodeSeq;
}
else {
return undefined;
}
}
// eslint-disable-next-line @typescript-eslint/no-unused-vars
canApply(nodes) {
return true;
}
}
//# sourceMappingURL=optimization.js.map