}

def XlaInlineDeviceOpsPass : Pass<"tf-xla-inline-device-ops", "ModuleOp"> {
  let summary = "Inline all Cluster op based in the parent region";
  let constructor = "TFDevice::CreateXlaInlineDeviceOpsPass()";
  let dependentDialects = ["tf_device::TensorFlowDeviceDialect"];
}

def XlaRewritePass : Pass<"tf-xla-rewrite", "mlir::ModuleOp"> {
  let summary = "Rewrites partition calls into Xla launch ops to make the attached function run on XLA.";

#include "tensorflow/compiler/mlir/tf2xla/transforms/passes.h"
#include "tensorflow/compiler/mlir/tf2xla/transforms/xla_legalize_targets.h"
#include "xla/mlir_hlo/mhlo/IR/hlo_ops.h"
#include "xla/mlir_hlo/mhlo/transforms/rewriters.h"
#include "xla/mlir_hlo/mhlo/utils/type_conversion.h"
#include "tensorflow/core/lib/monitoring/counter.h"

namespace mlir {
namespace mhlo {
namespace {

#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_traits.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/xla_sharding_util.h"
#include "xla/client/sharding_builder.h"
#include "xla/xla_data.pb.h"
#include "tsl/platform/errors.h"

namespace tensorflow {
namespace tf2xla {
namespace internal {
namespace {

#include "tensorflow/compiler/tf2xla/layout_util.h"
#include "tensorflow/compiler/tf2xla/xla_helpers.h"
#include "xla/client/client_library.h"
#include "xla/shape.h"
#include "xla/stream_executor/platform_manager.h"
#include "xla/translate/mhlo_to_hlo/type_to_shape.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/lib/monitoring/cell_reader.h"

      llvm::cl::desc("Choose target opset."),
      llvm::cl::values(
          clEnumValN(OpSet::TF, "TF",
                     "Uses TF ops that mimic quantization behavior"),
          clEnumValN(OpSet::XLA, "XLA", "Uses TF XLA ops"),
          clEnumValN(OpSet::UNIFORM_QUANTIZED, "UNIFORM_QUANTIZED",
                     "Uses TF Uniform Quantized ops"))};

  // Initialize for tests.
  void initializeForTest() {

#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/compiler/mlir/tf2xla/transforms/utils.h"
#include "xla/mlir_hlo/mhlo/IR/hlo_ops.h"
#include "xla/mlir_hlo/utils/convert_op_folder.h"
#include "xla/mlir_hlo/utils/hlo_utils.h"
#include "xla/xla_data.pb.h"

namespace mlir {
namespace mhlo {

namespace {

constexpr absl::string_view kGroupSizeAttrName =

#include "tensorflow/compiler/mlir/tensorflow/ir/tf_types.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/device_util.h"
#include "tensorflow/compiler/mlir/utils/string_container_utils.h"
#include "xla/array4d.h"
#include "xla/service/computation_placer.h"
#include "xla/xla_data.pb.h"
#include "tensorflow/core/framework/types.h"
#include "tensorflow/core/protobuf/tpu/topology.pb.h"
#include "tensorflow/core/util/device_name_utils.h"

#include "tensorflow/compiler/mlir/quantization/tensorflow/passes/passes.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/quantization_options.pb.h"
#include "tensorflow/compiler/mlir/tensorflow/transforms/passes.h"
#include "xla/mlir_hlo/mhlo/transforms/passes.h"

namespace tensorflow {
namespace quantization {
namespace {

using ::tensorflow::quantization::QuantizationOptions;

void AddConvertTpuToCpuModelPasses(mlir::OpPassManager &pm) {

      return errors::InvalidArgument(
          "Undeclared output of XLA computation. Some common causes of this "
          "error are: 1) variable initializers that depend on the XLA "
          "computation; 2) gradient computations that depend on the XLA "
          "computation, which can be mitigated by moving gradient computations "
          "inside XLA computation. Offending edge: ",

            "@local_xla//xla/service:hlo_profile_printer_data_cc",
        ] or []) + (include_standard_runtime_deps and [
            # TODO(cwhipkey): only depend on kernel code that the model actually
            # needed.
            "@local_xla//xla/service/cpu:runtime_conv2d",
            "@local_xla//xla/service/cpu:runtime_custom_call_status",
            "@local_xla//xla/service/cpu:runtime_key_value_sort",