using func::FuncOp;
using mlir::ModuleOp;

static constexpr char kMlirModuleStr[] = R"(
module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 1442 : i32}} {

    func.func @main(%arg0: tensor<f32>) {
      %elem = "tf.Const"() <{value = dense<0.000000e+00> : tensor<f32>}> : () -> tensor<f32>
      %shape = "tf.Const"() <{value = dense<> : tensor<0xi64>}> : () -> tensor<0xi64>
      %zero = "tf.Fill"(%shape, %elem) : (tensor<0xi64>, tensor<f32>) -> tensor<f32>
      tf_device.replicate([%arg0, %zero] as %arg1: tensor<f32>) {n = 2 : i32} {
        %arg1_id = "tf_device.launch"() <{device = "TPU_REPLICATED_HOST_0"}> ({

using ::xla::XlaComputation;
using ::xla::XlaOp;

static constexpr char kMlirModuleStr[] = R"(
module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 1442 : i32}} {

    SmallVector<Value> start_indices(rank);
    for (auto [i, start_index] : llvm::enumerate(op.getStartIndices())) {
      // Start indices should be casted from tensor<i64> to tensor<1xi64>.
      auto cast = rewriter.create<TFL::BitcastOp>(
          op->getLoc(), single_element_type, start_index);
      int64_t upper_limit_idx = operand_shape[i] - slice_sizes[i];
      auto upper_limit_attr =

// For example,
//
//   %0 = tf.Pack(%input) {axis = 0} // %input : tensor<2x3xf32>
//
// can be canonicalized to
//
//   %shape = "tf.Const"() {value = dense<[1, 2, 3]> : tensor<3xi64>}
//   %0 = tf.Reshape(%input, %shape)
struct ConvertPackToReshape : public OpRewritePattern<PackOp> {
  using OpRewritePattern<PackOp>::OpRewritePattern;

  LogicalResult matchAndRewrite(PackOp pack_op,