Function getSourceCoords

tfjs-backend-webgpu/src/tile_webgpu.ts:59–73

(rank: number, uniformPrefix = '')

}

function getSourceCoords(rank: number, uniformPrefix = ''): string {
  if (rank >= 5) {
    throw Error(`Tile for rank ${rank} is not yet supported`);
  }
  if (rank === 1) {
    return `(resRC % ${uniformPrefix}aShape)`;
  }

  const currentCoords = ['resRC.x', 'resRC.y', 'resRC.z', 'resRC.w'];
  const sourceCoords = [];
  for (let i = 0; i < rank; i++) {
    sourceCoords.push(`(${currentCoords[i]} % ${uniformPrefix}aShape[${i}])`);
  }
  return sourceCoords.join();
}