namespace mlir {
namespace mhlo {

ConstantOp GetScalarConstOfType(Type ty, Location loc, int64_t raw_value,
                                OpBuilder* builder) {
  return builder->create<ConstantOp>(loc, hlo::getScalarOfType(ty, raw_value));
}

ConstantOp GetScalarNegZeroOfType(Type ty, Location loc, OpBuilder* builder) {
  return builder->create<ConstantOp>(loc, hlo::getScalarNegZeroOfType(ty));
}

    func.func @main(%arg0: tensor<?xi32, #mhlo.type_extensions<bounds = [1]>>) -> tensor<?xi32, #mhlo.type_extensions<bounds = [1]>> {
      %0 = "tf.Identity"(%arg0) : (tensor<?xi32, #mhlo.type_extensions<bounds = [1]>>) -> tensor<?xi32, #mhlo.type_extensions<bounds = [1]>>
      func.return %0 : tensor<?xi32, #mhlo.type_extensions<bounds = [1]>>
    }
  })";

#define TENSORFLOW_COMPILER_MLIR_TF2XLA_TRANSFORMS_LEGALIZATION_OP_CONFIG_H_

#include "mlir/IR/Operation.h"  // from @llvm-project
#include "mlir/Support/TypeID.h"  // from @llvm-project

namespace mlir {
namespace mhlo {

// Given the Operation, return True if the op is legalized with MLIR.
// Returns false if it is legalized with an old TF2XLA Kernel.
bool IsOpLegalizedWithMlir(Operation& op);

}

// -----

func.func @illegal_mhlo_uniform_quantize(%arg0: tensor<?x?xf32>) -> tensor<?x?x!quant.uniform<i8:f32, 1.000000e+00:3>> {
  // expected-error@+1 {{'mhlo.uniform_quantize' op is illegal as it is a UQ op or contains uq/qint types}}
  %0 = mhlo.uniform_quantize %arg0 : (tensor<?x?xf32>) -> tensor<?x?x!quant.uniform<i8:f32, 1.000000e+00:3>>
  return %0 : tensor<?x?x!quant.uniform<i8:f32, 1.000000e+00:3>>
}

// -----

#include "mlir/Transforms/DialectConversion.h"  // from @llvm-project
#include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h"  // IWYU pragma: keep
#include "xla/mlir_hlo/mhlo/IR/hlo_ops.h"

namespace mlir {
namespace odml {

LogicalResult ConvertCustomCallOp::matchAndRewrite(
    mhlo::CustomCallOp mhlo_custom_call, OpAdaptor adaptor,
    ConversionPatternRewriter& rewriter) const {
  auto tfl_custom = rewriter.create<TFL::CustomOp>(

    func.return %4 : tensor<?xi32>
  }

  // CHECK-LABEL: func @_func
  // CHECK-SAME: %[[ARG0:.*]]: tensor<?xi32> {mhlo.is_same_data_across_replicas = true},
  // CHECK-SAME: %[[ARG1:.*]]: tensor<?xi32>,
  // CHECK-SAME: %[[ARG2:.*]]: tensor<!tf_type.resource<tensor<?xi32>>> {mhlo.is_same_data_across_replicas = true}

#include "tensorflow/compiler/mlir/tf2xla/transforms/test_utils.h"
#include "tsl/platform/statusor.h"

namespace tensorflow {
namespace tf2xla {
namespace internal {

namespace {

using mlir::mhlo::test::GetMlirModuleFromString;

class VerifyClusteringPassTest : public testing::Test {
 protected:
  void CreateModule(const char* module_string) {
    TF_ASSERT_OK_AND_ASSIGN(module_,

#include "tsl/platform/statusor.h"

namespace mlir {
namespace mhlo {
namespace test {

// Given a raw string, return a ModuleOp that can be used with the given
// MLIRContext.
absl::StatusOr<OwningOpRef<ModuleOp>> GetMlirModuleFromString(
    absl::string_view module_string, MLIRContext* mlir_context);

}  // namespace test
}  // namespace mhlo
}  // namespace mlir

namespace mlir {
namespace mhlo {
namespace {

mlir::DialectRegistry GetDefaultDialectRegistry() {
  mlir::DialectRegistry registry;

  registry.insert<arith::ArithDialect>();
  registry.insert<func::FuncDialect>();
  registry.insert<tensor::TensorDialect>();
  registry.insert<shape::ShapeDialect>();
  registry.insert<TF::TensorFlowDialect>();
  registry.insert<chlo::ChloDialect>();

  return registry;
}

#include "tensorflow/compiler/mlir/tf2xla/internal/clustering_bridge_passes.h"
#include "tensorflow/compiler/mlir/tf2xla/transforms/passes.h"
#include "xla/mlir_hlo/mhlo/transforms/passes.h"

namespace tensorflow {

void PopulateLowerToMlProgramAndHloPipeline(mlir::OpPassManager& pm) {
  tensorflow::tf2xla::internal::AddNonReplicatedBridgeClusteringPipelinePasses(
      pm);