@hoff97/tensor-js
Version:
PyTorch like deep learning inferrence library
55 lines • 1.67 kB
JavaScript
import { CPUTensor } from '../../tensor/cpu/tensor';
import { incrementIndex } from '../../util/shape';
export function pad(x, pads, mode, value) {
const rank = x.shape.length;
const resultShape = [...x.shape];
for (let i = 0; i < rank; i++) {
resultShape[i] += pads[i] + pads[i + rank];
}
const Y = new CPUTensor(resultShape, undefined, x.dtype);
const ix = new Array(rank).fill(0);
const inputIx = new Array(rank).fill(0);
for (let i = 0; i < Y.size; i++) {
let allInRange = true;
for (let j = 0; j < rank; j++) {
inputIx[j] = ix[j] - pads[j];
if (inputIx[j] < 0 || inputIx[j] >= x.shape[j]) {
allInRange = false;
}
}
Y.set(i, getPadValue(x, inputIx, mode, value, allInRange));
incrementIndex(ix, resultShape);
}
return Y;
}
function getPadValue(x, index, mode, value, allInRange) {
if (allInRange) {
return x.get(index);
}
const rank = x.shape.length;
if (mode === 'constant') {
return value;
}
else if (mode === 'edge') {
for (let j = 0; j < rank; j++) {
if (index[j] < 0) {
index[j] = 0;
}
else if (index[j] >= x.shape[j]) {
index[j] = x.shape[j] - 1;
}
}
}
else {
for (let j = 0; j < rank; j++) {
if (index[j] < 0) {
index[j] = -index[j];
}
else if (index[j] >= x.shape[j]) {
index[j] = 2 * x.shape[j] - index[j] - 2;
}
}
}
return x.get(index);
}
//# sourceMappingURL=pad.js.map