Function getCoordsFromIndexSnippet

tfjs-backend-webgpu/src/webgpu_program.ts:450–499

(shape: number[], name = '')

 * dimension coordinate.
 */
export function getCoordsFromIndexSnippet(shape: number[], name = ''): string {
  const rank = shape.length;
  const funcName = name !== '' ?
      `get${name.charAt(0).toUpperCase() + name.slice(1)}CoordsFromIndex` :
      'getCoordsFromIndex';
  const stridesName = name !== '' ?
      `${name.charAt(0).toLowerCase() + name.slice(1)}ShapeStrides` :
      `outShapeStrides`;

  if (rank <= 1) {
    return `fn ${funcName}(index : i32) -> i32 { return index; }`;
  }

  const strides = util.computeStrides(shape);
  const dtype = getCoordsDataType(rank);

  const coords: string[] = [];
  for (let i = 0; i < rank; i++) {
    coords.push(`d${i}`);
  }

  if (strides.length === 1) {
    return `    fn ${funcName}(index : i32) -> vec2<i32> {
      let d0 = index / uniforms.${
        stridesName}; let d1 = index - d0 * uniforms.${stridesName};
      return vec2<i32>(d0, d1);
    }`;
  }
  let snippet;
  snippet = 'var index2 = index;' +
      strides
          .map((_, i) => {
            const line1 = `let ${coords[i]} = index2 / uniforms.${
                stridesName}.${getCoordsXYZ(i)}`;
            const line2 = i === strides.length - 1 ?
                `let ${coords[i + 1]} = index2 - ${coords[i]} * uniforms.${
                    stridesName}.${getCoordsXYZ(i)}` :
                `index2 = index2 - ${coords[i]} * uniforms.${stridesName}.${
                    getCoordsXYZ(i)}`;
            return `${line1}; ${line2};`;
          })
          .join('');

  return `
    fn ${funcName}(index : i32) -> ${dtype} {
      ${snippet}
      return ${dtype}(${coords.join(',')});
    }
  `;
}

function getInputAtCoordsSnippet(
    inputInfo: InputInfo, component: number): string {