@tensorflow/tfjs-core
Version:
Hardware-accelerated JavaScript library for machine intelligence
49 lines • 1.53 kB
JavaScript
export function getBroadcastDims(inShape, outShape) {
var inRank = inShape.length;
var dims = [];
for (var i = 0; i < inRank; i++) {
var dim = inRank - 1 - i;
var a = inShape[dim] || 1;
var b = outShape[outShape.length - 1 - i] || 1;
if (b > 1 && a === 1) {
dims.unshift(dim);
}
}
return dims;
}
export function getReductionAxes(inShape, outShape) {
var result = [];
for (var i = 0; i < outShape.length; i++) {
var inDim = inShape[inShape.length - i - 1];
var outAxis = outShape.length - i - 1;
var outDim = outShape[outAxis];
if (inDim == null || (inDim === 1 && outDim > 1)) {
result.unshift(outAxis);
}
}
return result;
}
export function broadcastDimsAreOuter(dims) {
for (var i = 0; i < dims.length; i++) {
if (dims[i] !== i) {
return false;
}
}
return true;
}
export function assertAndGetBroadcastShape(shapeA, shapeB) {
var result = [];
var errMsg = "Operands could not be broadcast together with shapes " +
(shapeA + " and " + shapeB + ".");
var l = Math.max(shapeA.length, shapeB.length);
for (var i = 0; i < l; i++) {
var a = shapeA[shapeA.length - i - 1] || 1;
var b = shapeB[shapeB.length - i - 1] || 1;
if (a > 1 && b > 1 && a !== b) {
throw Error(errMsg);
}
result.unshift(Math.max(a, b));
}
return result;
}
//# sourceMappingURL=broadcast_util.js.map