Function kernelFunc

tfjs-backend-wasm/src/kernels/argminmax_kernel.ts:40–78

(args: {
    backend: BackendWasm,
    inputs: ArgMinInputs&ArgMaxInputs,
    attrs: ArgMinAttrs&ArgMaxAttrs,
  })

  }

  function kernelFunc(args: {
    backend: BackendWasm,
    inputs: ArgMinInputs&ArgMaxInputs,
    attrs: ArgMinAttrs&ArgMaxAttrs,
  }): TensorInfo {
    const {backend, inputs, attrs} = args;
    const {axis} = attrs;
    const {x} = inputs;
    const xId = backend.dataIdMap.get(x.dataId).id;
    let inputId = xId;
    let input = x;

    const {transposed, axes, inputWasTransposed} =
        permuteAxesAndTranspose(x, axis, backend);

    if (inputWasTransposed) {
      const transposedId = backend.dataIdMap.get(transposed.dataId).id;
      if (transposedId !== xId) {
        // transpose was not a no-op. We will need to dispose of this
        // once we are done.
        input = transposed;
        inputId = transposedId;
      }
    }

    const outShape = input.shape.slice(0, -1);
    const out = backend.makeOutput(outShape, 'int32');
    const outId = backend.dataIdMap.get(out.dataId).id;
    const outerSize = util.sizeFromShape(out.shape);
    const innerSize = input.shape[axes[0]];
    wasmFunc(inputId, CppDType[input.dtype], outerSize, innerSize, outId);

    if (inputWasTransposed) {
      // dispose of the transposed tensor.
      backend.disposeData(transposed.dataId);
    }

    return out;
  }

  return {
    kernelName,