Function searchSorted

tfjs-backend-webgpu/src/kernels/SearchSorted.ts:23–37

(args: {
  inputs: SearchSortedInputs,
  backend: WebGPUBackend,
  attrs: SearchSortedAttrs
})

import {SearchSortedProgram} from '../search_sorted_webgpu';

export function searchSorted(args: {
  inputs: SearchSortedInputs,
  backend: WebGPUBackend,
  attrs: SearchSortedAttrs
}): TensorInfo {
  const {inputs, backend, attrs} = args;
  const {sortedSequence, values} = inputs;
  const {side} = attrs;

  const program =
      new SearchSortedProgram([values.shape[0], values.shape[1]], side);
  const uniformData = [{type: 'int32', data: [sortedSequence.shape[1]]}];
  return backend.runWebGPUProgram(
      program, [sortedSequence, values], 'int32', uniformData);
}

export const searchSortedConfig: KernelConfig = {
  kernelName: SearchSorted,