func.func @sigmoid(%arg0: tensor<2xf32>) -> tensor<2xf32> {
  %0 = mhlo.constant dense<5.000000e-01> : tensor<f32>
  %1 = mhlo.constant dense<2> : tensor<1xi64>
  %2 = mhlo.constant dense<5.000000e-01> : tensor<2xf32>
  %3 = mhlo.multiply %arg0, %2 : tensor<2xf32>
  %4 = "mhlo.tanh"(%3) : (tensor<2xf32>) -> tensor<2xf32>
  %5 = mhlo.multiply %4, %2 : tensor<2xf32>
  %6 = mhlo.add %5, %2 : tensor<2xf32>
  func.return %6 : tensor<2xf32>
}

// CHECK:           [[LHSRE:%.*]] = mhlo.real [[LHS]]
// CHECK:           [[LHSIM:%.*]] = mhlo.imag [[LHS]]
// CHECK:           [[LHSIMNEG:%.*]] = mhlo.negate [[LHSIM]]
// CHECK:           [[LHSCONJ:%.*]] = mhlo.complex [[LHSRE]], [[LHSIMNEG]]
// CHECK:           [[RHSRE:%.*]] = mhlo.real [[RHS]]
// CHECK:           [[RHSIM:%.*]] = mhlo.imag [[RHS]]
// CHECK:           [[RHSIMNEG:%.*]] = mhlo.negate [[RHSIM]]

#include "tensorflow/compiler/mlir/quantization/common/attrs_and_constraints.h"
#include "xla/mlir_hlo/mhlo/IR/hlo_ops.h"

namespace mlir {
namespace odml {

class FuseMhloMulAndConvolutionPattern : public OpRewritePattern<mhlo::MulOp> {
 public:
  using OpRewritePattern<mhlo::MulOp>::OpRewritePattern;

  LogicalResult matchAndRewrite(mhlo::MulOp mul_op,
                                PatternRewriter &rewriter) const override {

// CHECK-LABEL: @unfold_splat_constant_float
func.func @unfold_splat_constant_float() -> tensor<1x750xf32> {
  %cst = mhlo.constant dense<7.680000e+02> : tensor<1x750xf32>
  func.return %cst : tensor<1x750xf32>

  // CHECK-DAG: %0 = mhlo.constant dense<7.680000e+02> : tensor<f32>
  // CHECK: %1 = "mhlo.broadcast_in_dim"(%0) <{broadcast_dimensions = dense<> : tensor<0xi64>}> : (tensor<f32>) -> tensor<1x750xf32>

  // Compiles the given Operation with XlaBuilder and imports the generated HLO
  // via the HLO -> MHLO importer.
  absl::StatusOr<mhlo::TupleOp> CompileWithHloImporter(
      tensorflow::OpKernelContext& op_context);

  // Import the given XlaComputation into the parent module. Returns the given
  // generated function.
  absl::StatusOr<mhlo::TupleOp> ImportXlaComputation(
      xla::XlaComputation& computation);

  return b.createOrFold<mhlo::DynamicBroadcastInDimOp>(
      broadcast_to_type, epsilon, shape_value, dims);
}

class UnfuseBatchNormTrainingPattern
    : public OpRewritePattern<mhlo::BatchNormTrainingOp> {
 public:
  using OpRewritePattern<mhlo::BatchNormTrainingOp>::OpRewritePattern;

  LogicalResult matchAndRewrite(mhlo::BatchNormTrainingOp bn_op,

  %0 = mhlo.constant dense<1.000000e+00> : tensor<f64>
  %1 = mhlo.convert %0 : (tensor<f64>) -> tensor<i64>
  %2 = mhlo.reshape %1 : (tensor<i64>) -> tensor<1xi64>
  // expected-error @+1 {{Node `mhlo.dynamic_iota` must have compile-time constant}}
  %3 = "mhlo.dynamic_iota"(%2) <{iota_dimension = 0 : i64}> : (tensor<1xi64>) -> tensor<?xi32>
  %4 = mhlo.multiply %3, %3 : tensor<?xi32>
  return %4 : tensor<?xi32>

    // 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 = [],