// Shape of the final output. (Except for dimension folding in the
    // single diagonal case.)
    Shape output_shape;
    for (int i = 0; i < num_dims - 2; i++) {
      output_shape.push_back(input_type.getDimSize(i));
    }
    output_shape.push_back(num_diags);
    output_shape.push_back(max_diag_len);

    // A slice is the shape of what GatherOp copies per lookup. So the last