Function select

tfjs-backend-webgpu/src/kernels/Select.ts:23–32

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

import {SelectProgram} from '../select_webgpu';

export function select(args: {inputs: SelectInputs, backend: WebGPUBackend}):
    TensorInfo {
  const {inputs, backend} = args;
  const {condition, t, e} = inputs;

  const program =
      new SelectProgram(condition.shape.length, t.shape, t.shape.length);
  return backend.runWebGPUProgram(
      program, [condition, t, e], upcastType(t.dtype, e.dtype));
}

export const selectConfig: KernelConfig = {
  kernelName: Select,