Function assertAndGetBroadcastShape

tfjs-core/src/ops/broadcast_util.ts:60–87

(
    shapeA: number[], shapeB: number[])

}

export function assertAndGetBroadcastShape(
    shapeA: number[], shapeB: number[]): number[] {
  const l = Math.max(shapeA.length, shapeB.length);
  const result = new Array(l);

  for (let i = 0; i < l; i++) {
    let a = shapeA[shapeA.length - i - 1];
    if (a == null) {
      a = 1;
    }
    let b = shapeB[shapeB.length - i - 1];
    if (b == null) {
      b = 1;
    }
    if (a === 1) {
      result[l - i - 1] = b;
    } else if (b === 1) {
      result[l - i - 1] = a;
    } else if (a !== b) {
      const errMsg = `Operands could not be broadcast together with shapes ` +
          `${shapeA} and ${shapeB}.`;
      throw Error(errMsg);
    } else {
      result[l - i - 1] = a;
    }
  }
  return result;
}