Function getInputAtCoordsSnippet

tfjs-backend-webgpu/src/webgpu_program.ts:501–532

(
    inputInfo: InputInfo, component: number)

}

function getInputAtCoordsSnippet(
    inputInfo: InputInfo, component: number): string {
  const texName = inputInfo.name;
  const rank = inputInfo.shape.length;
  const type = getCoordsDataType(rank);
  const funcName = 'get' + texName.charAt(0).toUpperCase() + texName.slice(1);
  const dims = ['d0', 'd1', 'd2', 'd3', 'd4', 'd5'].slice(0, rank);
  const inputs = dims.map(d => `${d} : i32`).join(', ');

  if (rank < 1) {
    return `
      fn ${funcName}() -> ${typeSnippet(component)} {
        return ${typeSnippet(component)}(${texName}[0]);
      }
    `;
  }

  const shapeStr =
      `uniforms.${texName.charAt(0).toLowerCase() + texName.slice(1)}Shape`;
  let rankStr = `${rank}D`;
  if (rank === 0) {
    rankStr = '1D';
  }

  return `
    fn ${funcName}(${inputs}) -> ${typeSnippet(component)} {
      return ${typeSnippet(component)}(${texName}[getIndexFromCoords${
      rankStr}(${type}(${dims.join(',')}),
        ${shapeStr})${component === 1 ? '' : ` / ${component}`}]);
    }
   `;
}

function getInputByOutputSnippet(
    inputInfo: InputInfo, outShape: number[], component: number,