rewriter.create<arith::ConstantOp>(op->getLoc(), type, begin_attr);
    auto end_attr = DenseElementsAttr::get<int32_t>(type, padded_end);
    auto end_op =
        rewriter.create<arith::ConstantOp>(op->getLoc(), type, end_attr);
    auto stride_attr = DenseElementsAttr::get<int32_t>(type, padded_stride);
    auto stride_op =
        rewriter.create<arith::ConstantOp>(op->getLoc(), type, stride_attr);

  OpBuilder builder(func.getBody());
  std::string empty_option_buffer;
  auto op = builder.create<CustomOp>(
      func.getLoc(), func.getFunctionType().getResults(), func.getArguments(),
      api, CustomOption(&builder, empty_option_buffer));
  builder.create<func::ReturnOp>(func.getLoc(), op.getResults());
  return success();
}

LogicalResult VerifyNgrams(func::FuncOp func) {

  assert(caller_region.empty() &&
         "Expected empty region for newly created ops");
  OpBuilder builder(caller_region);
  Block* entry = builder.createBlock(&caller_region);

  auto loc = op->getLoc();
  if (use_region_args) {
    auto inputs = func.getFunctionType().getInputs();
    entry->addArguments(inputs, SmallVector<Location>(inputs.size(), loc));
    args = entry->getArguments();
  }

        scf_yield_input.push_back(rewriter.create<CastOp>(
            current_terminator->getLoc(), std::get<0>(it), std::get<1>(it)));
      }

      rewriter.replaceOpWithNewOp<scf::YieldOp>(current_terminator,
                                                scf_yield_input);
    };

    Location loc = op.getLoc();

    // The condition of an `scf.if` op is a 1-bit signless integer. Whereas, the

  auto reshape_sizes = DenseIntElementsAttr::get(reshape_type, new_shape);
  auto reshape_value =
      builder->create<TF::ConstOp>(input.getLoc(), reshape_sizes);
  return builder->create<TF::ReshapeOp>(input.getLoc(), reshape_result_type,
                                        input, reshape_value);
}

// Creates transpose op for shape to depth transform.

        if (IsCompatibleTypeWithTFLCastOp(value.getType()) &&
            IsCompatibleTypeWithTFLCastOp(type)) {
          auto cast = b.create<CastOp>(yield_op->getLoc(), type, value);
          args.push_back(cast);
        } else {
          auto cast = b.create<TF::CastOp>(yield_op->getLoc(), type, value);
          args.push_back(cast);
        }
      }
    }
    args.append(new_args.begin(), new_args.end());
  } else {

    auto var_handle_op =
        builder.create<TF::VarHandleOp>(const_op.getLoc(),
                                        /*resource=*/resource_type,
                                        /*container=*/"", shared_name);

    auto read_variable_op = builder.create<TF::ReadVariableOp>(
        const_op.getLoc(), const_op.getType(), var_handle_op);

        mlir::cast<TensorType>(input.getType()), kNchwToNhwcPermutation);

    auto input_transpose_op = rewriter.create<mlir::stablehlo::TransposeOp>(
        op.getLoc(), /*resultType0=*/new_input_tensor_type, /*operand=*/input,
        rewriter.getDenseI64ArrayAttr(kNchwToNhwcPermutation));

    // Transpose the filter tensor: [o, i, 0, 1] => [0, 1, i, o]
    Value filter = op->getOperand(1);

      return InlineRegion(if_op.getLoc(), rewriter, if_op,
                          &if_op.getThenRegion());
    }

    // Else branch
    if (matchPattern(if_op.getCondition(), m_Zero())) {
      if (if_op.getElseRegion().empty()) {
        // Remove the op
        rewriter.eraseOp(if_op);
        return success();
      } else {
        return InlineRegion(if_op.getLoc(), rewriter, if_op,

  }

  auto compile_op = builder->create<TF::_TPUCompileMlirOp>(
      cluster_func.getLoc(),
      /*compilation_status=*/compilation_status_type, /*program=*/
      llvm::SmallVector<Type, 8>(num_cores_per_replica, program_type),
      compile_op_operands, txt_module, txt_metadata);

  return tensorflow::WrapOpInLaunch(builder, compile_op.getLoc(), compile_op,
                                    compilation_device);
}