namespace mlir {
namespace odml {
// Converts mhlo.dot_general to tfl.BatchMatMul. Reshape and Transpose ops will
// be inserted to convert to well-formed matrix multiply; i.e., mhlo.dot_general
// -> tfl.batch_matmul(mhlo.transpose(mhlo.reshape(operand)), ...).
// Note:
// 1) Reshape/transpose are inserted because tfl.BatchMatMul requires
// size(contracting_dimensions) = 1 and size(output_dim) = 1, whereas
// mhlo.dot_general has no such restriction.

    // CHECK-NEXT: %5 = "mhlo.slice"(%1) <{limit_indices = dense<2> : tensor<1xi64>, start_indices = dense<1> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>}> : (tensor<4xi32>) -> tensor<1xi32>
    // CHECK-NEXT: %6 = mhlo.reshape %5 : (tensor<1xi32>) -> tensor<i32>
    // CHECK-NEXT: %7 = mhlo.convert %6 : tensor<i32>

  func.return %0 : tensor<4xf32>
}

// CHECK-LABEL: func @mod
// CHLO-LABEL: func @mod
func.func @mod(%arg0: tensor<4xf32>, %arg1: tensor<4xf32>) -> tensor<4xf32> {
  // CHECK-NEXT:  mhlo.remainder %arg0, %arg1 : tensor<4xf32>
  // CHLO: chlo.broadcast_remainder
  %0 = "tf.Mod"(%arg0, %arg1) : (tensor<4xf32>, tensor<4xf32>) -> tensor<4xf32>
  func.return %0 : tensor<4xf32>
}

  ) -> tensor<?x?x8xi32> {
  // CHECK: %[[dot:.*]] = "mhlo.dot_general"
  // CHECK: %[[combined:.*]] = chlo.broadcast_add %[[zp_offset:.*]], %[[bias:.*]]
  // CHECK: %[[result:.*]] = chlo.broadcast_add %[[dot]], %[[combined]]
  // CHECK: return %[[result]]
  %0 = "mhlo.dot_general" (%lhs, %rhs) {
    dot_dimension_numbers = #mhlo.dot<
      lhs_batching_dimensions = [],
      rhs_batching_dimensions = [],

  %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>

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();

  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: %[[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]]

// 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<

  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