36 lines
2.1 KiB
JavaScript
36 lines
2.1 KiB
JavaScript
// Exact CPU pad specialization: copy contiguous blocks instead of converting
|
|
// every scalar index to/from coordinates. Complex tensors use the original.
|
|
module.exports = function install(face) {
|
|
var backend = face.tf.backend(), original = backend.pad, enabled = false;
|
|
backend.pad = function (input, paddings, constantValue) {
|
|
if (!enabled || input.dtype === 'complex64' || !input.shape.length) return original.call(this, input, paddings, constantValue);
|
|
var rank = input.shape.length, outShape = [], lastPadded = -1;
|
|
for (var axis = 0; axis < rank; axis++) {
|
|
if (!paddings[axis] || paddings[axis].length !== 2) return original.call(this, input, paddings, constantValue);
|
|
var before = paddings[axis][0], after = paddings[axis][1];
|
|
if (before < 0 || after < 0 || Math.floor(before) !== before || Math.floor(after) !== after) return original.call(this, input, paddings, constantValue);
|
|
if (before || after) lastPadded = axis;
|
|
outShape.push(input.shape[axis] + before + after);
|
|
}
|
|
var output = face.tf.buffer(outShape, input.dtype), source = this.readSync(input.dataId), target = output.values;
|
|
if (constantValue !== 0) target.fill(constantValue);
|
|
var firstBlockAxis = lastPadded < 0 ? 0 : lastPadded;
|
|
var blockSize = 1, padOffset = 0;
|
|
for (var dimension = 0; dimension < rank; dimension++) {
|
|
padOffset += paddings[dimension][0] * (dimension === rank - 1 ? 1 : output.strides[dimension]);
|
|
if (dimension >= firstBlockAxis) blockSize *= input.shape[dimension];
|
|
}
|
|
var blocks = input.size / blockSize;
|
|
for (var block = 0; block < blocks; block++) {
|
|
var remainder = block, destination = padOffset;
|
|
for (var outer = firstBlockAxis - 1; outer >= 0; outer--) {
|
|
var coordinate = remainder % input.shape[outer]; remainder = Math.floor(remainder / input.shape[outer]);
|
|
destination += coordinate * output.strides[outer];
|
|
}
|
|
target.set(source.subarray(block * blockSize, (block + 1) * blockSize), destination);
|
|
}
|
|
return output.toTensor();
|
|
};
|
|
return { setEnabled: function (value) { enabled = value === true; } };
|
|
};
|