@hoff97/tensor-js
Version:
PyTorch like deep learning inferrence library
23 lines • 623 B
JavaScript
export function poolResultShape(inputShape, axes, keepDims) {
const resultShape = [];
const sumShape = [];
const ixMap = [];
for (let i = 0; i < inputShape.length; i++) {
if (!axes.includes(i)) {
resultShape.push(inputShape[i]);
ixMap.push(i);
}
else {
if (keepDims) {
resultShape.push(1);
ixMap.push(i);
}
sumShape.push(inputShape[i]);
}
}
if (resultShape.length === 0) {
resultShape.push(1);
}
return [resultShape, ixMap];
}
//# sourceMappingURL=pool.js.map