Function complex

tfjs-backend-webgpu/src/kernels/Complex.ts:31–46

(args: {inputs: ComplexInputs, backend: WebGPUBackend})

 * components by calling disposeData on each.
 */
export function complex(args: {inputs: ComplexInputs, backend: WebGPUBackend}):
    TensorInfo {
  const {inputs, backend} = args;
  const {real, imag} = inputs;

  const complexInfo = backend.makeTensorInfo(real.shape, 'complex64');
  const complex = backend.tensorMap.get(complexInfo.dataId);

  const realTensorInfo = identity({inputs: {x: real}, backend});

  const imagTensorInfo = identity({inputs: {x: imag}, backend});

  complex.complexTensorInfos = {real: realTensorInfo, imag: imagTensorInfo};

  return complexInfo;
}

export const complexConfig: KernelConfig = {
  kernelName: Complex,