Function createComplexBinaryKernelImpl

tfjs-backend-cpu/src/utils/binary_utils.ts:135–194

(op: ComplexBinaryOperation)

 * Supports broadcast.
 */
export function createComplexBinaryKernelImpl(op: ComplexBinaryOperation):
    ComplexBinaryKernelImpl {
  return (aShape: number[], bShape: number[], aRealVals: Float32Array,
          aImagVals: Float32Array, bRealVals: Float32Array,
          bImagVals: Float32Array): [TypedArray, TypedArray, number[]] => {
    const resultShape = backend_util.assertAndGetBroadcastShape(aShape, bShape);
    const resultSize = util.sizeFromShape(resultShape);
    const resultRank = resultShape.length;
    const resultStrides = util.computeStrides(resultShape);

    const resultRealVals = util.getTypedArrayFromDType('float32', resultSize);
    const resultImagVals = util.getTypedArrayFromDType('float32', resultSize);

    const aBroadcastDims = backend_util.getBroadcastDims(aShape, resultShape);
    const bBroadcastDims = backend_util.getBroadcastDims(bShape, resultShape);

    const aVals = backend_util.mergeRealAndImagArrays(aRealVals, aImagVals);
    const bVals = backend_util.mergeRealAndImagArrays(bRealVals, bImagVals);

    const aRank = aShape.length;
    const aStrides = util.computeStrides(aShape);

    const bRank = bShape.length;
    const bStrides = util.computeStrides(bShape);

    if (aBroadcastDims.length + bBroadcastDims.length === 0) {
      for (let i = 0; i < resultRealVals.length; i++) {
        const aIdx = i % aVals.length;
        const bIdx = i % bVals.length;

        const result =
            op(aVals[aIdx * 2], aVals[aIdx * 2 + 1], bVals[bIdx * 2],
               bVals[bIdx * 2 + 1]);

        resultRealVals[i] = result.real;
        resultImagVals[i] = result.imag;
      }
    } else {
      for (let i = 0; i < resultRealVals.length; i++) {
        const loc = util.indexToLoc(i, resultRank, resultStrides);

        const aLoc = loc.slice(-aRank);
        aBroadcastDims.forEach(d => aLoc[d] = 0);
        const aIndex = util.locToIndex(aLoc, aRank, aStrides);

        const bLoc = loc.slice(-bRank);
        bBroadcastDims.forEach(d => bLoc[d] = 0);
        const bIndex = util.locToIndex(bLoc, bRank, bStrides);

        const opResult =
            op(aVals[aIndex * 2], aVals[aIndex * 2 + 1], bVals[bIndex * 2],
               bVals[bIndex * 2 + 1]);

        resultRealVals[i] = opResult.real;
        resultImagVals[i] = opResult.imag;
      }
    }
    return [resultRealVals, resultImagVals, resultShape];