constructor(shape){
        this.variableNames = [
            "x",
            "indices"
        ];
        // |n| Size of the original input of TopK.
        // |firstPass|indicates if this is the first time swap is being used which
        // means no indices input containing the top K is present yet.
        // |inc| Swaps pairs of indices (0, inc), (1, inc + 1), (2, inc + 2) ...
        this.customUniforms = [
            {
                name: "n",
                type: "int"
            },
            {
                name: "firstPass",
                type: "int"
            },
            {
                name: "negativeInf",
                type: "float"
            },
            {
                name: "dir",
                type: "int"
            },
            {
                name: "inc",
                type: "int"
            }
        ];
        this.outputShape = shape;
        this.userCode = `
       void main() {
         ivec2 coords = getOutputCoords();
         int batch = coords[0];
         int elemIdx = coords[1];

         // We compare elements pair-wise within a group of size 2 * inc.
         // The comparing rule for each group alternates between ascending
         // and descending. Within each group, we compare each pair at
         // positions i and i+inc. To decide whether an element at position i
         // is x0 or x1, we mod it by 2 * inc, if the result is smaller than
         // inc, it is in the first half of the group, we denote it as x0,
         // otherwise we denote it as x1.
         // For example, as shown in the Bitonic top K paper referenced above,
         // Figure5(a) shows that element[1] is in the
         // second half of the group when group size is 2, but it is in the
         // first half of the group when group size is 4.

         bool isFirstInPair = imod(elemIdx, 2 * inc) < inc;
         int i = isFirstInPair ? elemIdx : elemIdx - inc;

         int i0 = firstPass == 1 ? i : int(getIndices(batch, i));
         int i1 = firstPass == 1 ? i + inc : int(getIndices(batch, i + inc));
         float x0 = i0 < n ? getX(batch, i0) : negativeInf;
         float x1 = i1 < n ? getX(batch, i1) : negativeInf;

         // Denotes which direction indices are in (ascending or descending).
         bool reverse = imod(elemIdx, 2 * dir) >= dir;
         bool isGreater = x0 > x1 || (x0 == x1 && i1 > i0);
         if (reverse == isGreater) { // Elements in opposite order of direction
           int iTemp = i0;
           i0 = i1;
           i1 = iTemp;
         }
         if (isFirstInPair) {
            setOutput(float(i0));
         } else {
            setOutput(float(i1));
         }
       }
     `;
    }
@param shape desired output shape (can be larger than input shape, output
                                   will be padded with -Infinity)