squeeze_dims.push_back(squeeze_dim);
    }

    // Collect the unsqueezed dimensions.
    llvm::SmallVector<Value> dims;
    for (int64_t i = 0; i != input_rank; ++i) {
      if (llvm::is_contained(squeeze_dims, i)) continue;
      dims.push_back(rewriter.create<tensor::DimOp>(op.getLoc(), input, i));
    }

    auto from_extents =
        rewriter.create<tensor::FromElementsOp>(op.getLoc(), dims);