Function concat

tfjs-backend-webgpu/src/kernels/Concat.ts:25–49

(
    args: {inputs: ConcatInputs, attrs: ConcatAttrs, backend: WebGPUBackend})

import {identity} from './Identity';

export function concat(
    args: {inputs: ConcatInputs, attrs: ConcatAttrs, backend: WebGPUBackend}):
    TensorInfo {
  const {inputs, backend, attrs} = args;
  const {axis} = attrs;

  const $axis = util.parseAxisParam(axis, inputs[0].shape)[0];

  const shapes = inputs.map(t => t.shape);
  backend_util.assertParamsConsistent(shapes, $axis);

  const outShape =
      backend_util.computeOutShape(inputs.map(t => t.shape), $axis);
  if (util.sizeFromShape(outShape) === 0) {
    return backend.makeTensorInfo(outShape, inputs[0].dtype, []);
  }

  // Keep only non-empty tensors (ignore tensors with 0 in their shape).
  const $inputs = inputs.filter(t => util.sizeFromShape(t.shape) > 0);
  if ($inputs.length === 1) {
    return identity({inputs: {x: $inputs[0]}, backend});
  }

  return concatImpl($inputs, $axis, backend);
}

export const concatConfig: KernelConfig = {
  kernelName: Concat,