Function binaryKernelFunc

tfjs-backend-cpu/src/utils/binary_utils.ts:38–129

(
    name: string, simpleImpl: SimpleBinaryKernelImpl,
    complexImpl?: ComplexBinaryKernelImpl, dtype?: DataType)

 *     comparison kernels, such as Equal, Less, Greater, etc.
 */
export function binaryKernelFunc(
    name: string, simpleImpl: SimpleBinaryKernelImpl,
    complexImpl?: ComplexBinaryKernelImpl, dtype?: DataType): KernelFunc {
  if (complexImpl == null) {
    return ({inputs, backend}) => {
      const {a, b} = inputs as BinaryInputs;
      const cpuBackend = backend as MathBackendCPU;

      assertNotComplex([a, b], name);

      const aVals = cpuBackend.data.get(a.dataId).values as TypedArray;
      const bVals = cpuBackend.data.get(b.dataId).values as TypedArray;

      const decodedAVals = a.dtype === 'string' ?
          // tslint:disable-next-line: no-any
          backend_util.fromUint8ToStringArray(aVals as any as Uint8Array[]) :
          aVals;
      const decodedBVals = a.dtype === 'string' ?
          // tslint:disable-next-line: no-any
          backend_util.fromUint8ToStringArray(bVals as any as Uint8Array[]) :
          bVals;
      const $dtype = dtype || a.dtype;

      const [resultData, resultShape] =
          simpleImpl(a.shape, b.shape, decodedAVals, decodedBVals, $dtype);

      return cpuBackend.makeTensorInfo(resultShape, $dtype, resultData);
    };
  }

  return ({inputs, backend}) => {
    const {a, b} = inputs as BinaryInputs;
    const cpuBackend = backend as MathBackendCPU;

    if (a.dtype === 'complex64' || b.dtype === 'complex64') {
      const $aComplex = cast(
          {inputs: {x: a}, backend: cpuBackend, attrs: {dtype: 'complex64'}});

      const $aComplexVals = cpuBackend.data.get($aComplex.dataId);

      const aReal = $aComplexVals.complexTensorInfos.real;
      const aImag = $aComplexVals.complexTensorInfos.imag;

      const aRealVals =
          cpuBackend.data.get(aReal.dataId).values as Float32Array;
      const aImagVals =
          cpuBackend.data.get(aImag.dataId).values as Float32Array;

      const $bComplex = cast(
          {inputs: {x: b}, backend: cpuBackend, attrs: {dtype: 'complex64'}});

      const $bComplexVals = cpuBackend.data.get($bComplex.dataId);

      const bReal = $bComplexVals.complexTensorInfos.real;
      const bImag = $bComplexVals.complexTensorInfos.imag;

      const bRealVals =
          cpuBackend.data.get(bReal.dataId).values as Float32Array;