include "mlir/IR/OpBase.td"
include "mlir/IR/PatternBase.td"
include "mlir/Dialect/Func/IR/FuncOps.td"
include "mlir/Dialect/Arith/IR/ArithOps.td"
include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.td"
include "mlir/Dialect/Arith/IR/ArithOps.td"
include "tensorflow/compiler/mlir/lite/quantization/ir/QuantOps.td"
include "tensorflow/compiler/mlir/quantization/common/attrs_and_constraints.td"

#include "mlir/IR/PatternMatch.h"  // from @llvm-project
#include "mlir/Support/LLVM.h"  // from @llvm-project
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"  // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"

namespace mlir {
namespace quant {

// Copied from tensorflow/compiler/mlir/lite/transforms/prepare_tf.cc.
// By removing identity ops, constant operands with dynamic shapes have static

include "mlir/Dialect/Func/IR/FuncOps.td"
include "mlir/Dialect/Arith/IR/ArithOps.td"
include "tensorflow/compiler/mlir/quantization/common/attrs_and_constraints.td"
include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.td"

//===----------------------------------------------------------------------===//
// Pattern rules for converting bfloat16 operations to fp32 conversions.

#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
#include "mlir/Interfaces/SideEffectInterfaces.h"  // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"

namespace tensorflow {
namespace tfrt_compiler {

bool IsTensorArrayOp(mlir::Operation* op) {
  return llvm::isa<mlir::TF::TensorArrayV3Op, mlir::TF::TensorArrayScatterV3Op,

#include "mlir/IR/Dialect.h"  // from @llvm-project
#include "mlir/IR/Types.h"  // from @llvm-project
#include "mlir/Support/LogicalResult.h"  // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/compiler/mlir/tensorflow/transforms/constant_fold.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/convert_tensor.h"
#include "xla/stream_executor/stream_executor.h"

#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
#include "mlir/IR/Operation.h"  // from @llvm-project
#include "mlir/Pass/Pass.h"  // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/core/lib/monitoring/counter.h"

namespace mlir {
namespace tf2xla {
namespace internal {

auto* has_tpu_partitioned_call_streamz =
    tensorflow::monitoring::Counter<1>::New(

limitations under the License.
==============================================================================*/

include "mlir/IR/OpBase.td"
include "mlir/IR/PatternBase.td"
include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.td"
include "mhlo/IR/hlo_ops.td"
include "stablehlo/dialect/ChloOps.td"

def IsDefinedByConvOrDotGeneralOp : Constraint<Or<[
  CPred<"$0.getDefiningOp<mhlo::ConvolutionOp>()">,

  const auto cc_file_path = io::JoinPath(tmpdir, "test.cc");
  const auto internal_h_file_path = io::JoinPath(tmpdir, "test_internal.h");
  const auto internal_cc_file_path = io::JoinPath(tmpdir, "test_internal.cc");

  cc_op::WriteCCOps(ops, api_def_map, h_file_path, cc_file_path);

  TF_ASSERT_OK(ReadFileToString(env, h_file_path, h_file_text));
  TF_ASSERT_OK(
      ReadFileToString(env, internal_h_file_path, internal_h_file_text));
}

#include "tensorflow/compiler/mlir/init_mlir.h"
#include "tensorflow/compiler/mlir/lite/quantization/ir/QuantOps.h"
#include "tensorflow/compiler/mlir/tensorflow/dialect_registration.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/compiler/mlir/tfr/ir/tfr_ops.h"

int main(int argc, char **argv) {
  tensorflow::InitMlir y(&argc, &argv);

  mlir::registerAllPasses();

  mlir::DialectRegistry registry;

#include "mlir/Dialect/Traits.h"  // from @llvm-project
#include "tensorflow/compiler/mlir/quantization/tensorflow/quantization_options.pb.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"

namespace mlir {
namespace quant {

std::optional<TF::PartitionedCallOp> ApplyUniformQuantization(
    PatternRewriter& rewriter, TF::ConstOp op,