namespace mlir {
namespace odml {

// Creates a pass which unfuses MHLO batch norm inference op into arithmetic
// ops.
std::unique_ptr<Pass> createUnfuseBatchNormPass();

// Creates a pass which constant folds broadcast_in_dim op conditionally.
std::unique_ptr<Pass> createFoldBroadcastPass();

// Creates a pass which fuses MHLO binary element-wise ops and convolution op.
std::unique_ptr<Pass> createFuseConvolutionPass();

  %8 = mhlo.divide %7, %2 : tensor<64xf32>
  %9 = mhlo.floor %8 : tensor<64xf32>
  %10 = mhlo.convert %9 : (tensor<64xf32>) -> tensor<64xi32>
  %11 = mhlo.compare  LT, %10, %1,  SIGNED : (tensor<64xi32>, tensor<64xi32>) -> tensor<64xi1>
  %12 = mhlo.add %10, %0 : tensor<64xi32>
  %13 = mhlo.select %11, %12, %10 : tensor<64xi1>, tensor<64xi32>
  %14 = mhlo.reshape %13 : (tensor<64xi32>) -> tensor<64x1xi32>

  func.return
}

// CHECK-LABEL: func @func_with_sharding
// CHECK-SAME: (%{{[a-z0-9]+}}: tensor<*xi32> {mhlo.sharding = "\01\02\03"}, %{{[a-z0-9]+}}: tensor<*xi1> {mhlo.sharding = "\04\05\06"})
// CHECK-SAME: -> (tensor<*xi32> {mhlo.sharding = "\0A\0B\0C"}, tensor<*xi1> {mhlo.sharding = "\0D\0E\0F"})
func.func @func_with_sharding(%arg0: tensor<*xi32>, %arg1: tensor<*xi1>) -> (tensor<*xi32>, tensor<*xi1>) {

// CHECK-LABEL: testDotToDotGeneralVectorVector
func.func @testDotToDotGeneralVectorVector(%arg0: tensor<3072xf32>, %arg1: tensor<3072xf32>) -> tensor<f32> {
  %0 = "mhlo.dot"(%arg0, %arg1) : (tensor<3072xf32>, tensor<3072xf32>) -> tensor<f32>
  func.return %0 : tensor<f32>

// CHECK:      %[[RES:.*]] = "mhlo.dot_general"(%arg0, %arg1) <{
// CHECK-SAME:   dot_dimension_numbers = #mhlo.dot<

  // CHECK: %[[GROUP_SIZE:.*]] = mhlo.constant dense<2.000000e+00>
  // CHECK: %[[REDUCE:.*]] = "mhlo.all_reduce"
  // CHECK-SAME{LITERAL}: replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>
  // CHECK: mhlo.add
  // CHECK: mhlo.return
  // CHECK: %[[RESULT:.*]] = mhlo.divide %[[REDUCE]], %[[GROUP_SIZE]]
  // CHECK-NEXT: return %[[RESULT]]

  return %14 : tensor<32x10xf32>
}
// CHECK: ReadVariableOp
// CHECK: mhlo.convolution
// CHECK: mhlo.add
// CHECK: mhlo.maximum
// CHECK: mhlo.reshape
// CHECK: mhlo.dot
// CHECK: mhlo.add
// CHECK: mhlo.maximum
// CHECK: return{{.*}}tensor<32x10xf32>

// -----

// CHECK-LABEL: @mhlo_while
func.func private @mhlo_while() {
  // CHECK-NEXT: mhlo.constant

  // CHECK-DAG: %[[EPS_BCAST:.+]] = mhlo.constant dense<1.001000e-05> : tensor<256xf32>
  // CHECK-DAG: %[[VARIANCE_EPS:.+]] = mhlo.add %[[VARIANCE]], %[[EPS_BCAST]] : tensor<256xf32>
  // CHECK-DAG: %[[VARIANCE_EPS_RSQRT:.+]] = mhlo.rsqrt %[[VARIANCE_EPS]] : tensor<256xf32>
  // CHECK-DAG: %[[MULTIPLIER:.+]] = mhlo.multiply %[[VARIANCE_EPS_RSQRT]], %[[SCALE]] : tensor<256xf32>

  let summary = "Convert TF Quant ops to MHLO quantizated ops.";

  let description = [{
    Convert TF Quant ops to MHLO quant ops.
  }];

  let constructor = "mlir::quant::stablehlo::CreateConvertTFQuantOpsToMHLOPass()";
  let dependentDialects = ["TF::TensorFlowDialect", "chlo::ChloDialect",
                           "mhlo::MhloDialect", "tf_type::TFTypeDialect",

    return failure();

  auto compare_direction =
      is_argmax ? mhlo::ComparisonDirection::GT : mhlo::ComparisonDirection::LT;
  if (is_float) {
    mhlo::OrOp value_or = llvm::dyn_cast_or_null<mhlo::OrOp>(
        value_select.getOperand(0).getDefiningOp());
    if (!value_or) return failure();

    mhlo::CompareOp value_gt = llvm::dyn_cast_or_null<mhlo::CompareOp>(
        value_or.getLhs().getDefiningOp());

  // CHECK: [[TOKEN:%.*]] = mhlo.create_token : !mhlo.token
  %0 = "mhlo.create_token"() : () -> !mhlo.token

  // CHECK:               [[INFEED:%.*]]:3 = "mhlo.infeed"([[TOKEN]]) <{
  // CHECK-SAME{LITERAL}:   infeed_config = "", layout = [[1, 3, 2, 0], [1, 2, 0]]
  // CHECK-SAME:          }> : (!mhlo.token) -> (tensor<1x8x4x4xi32>, tensor<1x100x1xf32>, !mhlo.token)