Method getUserCode

tfjs-backend-webgpu/src/onehot_webgpu.ts:39–51

()

  }

  getUserCode(): string {
    const userCode = `
      ${main('index')} {
        if(index < uniforms.size) {
          let coords = getCoordsFromIndex(index);
          setOutputAtIndex(index, mix(uniforms.offValue, uniforms.onValue,
                                      f32(i32(round(getX(coords.x))) == coords.y)));
        }
      }
    `;

    return userCode;
  }
}