// CHECK-NEXT: %[[cmul:.*]] = mhlo.convert %[[mul]] : tensor<8x8x8x8xf32>
  // CHECK-NEXT: %[[init:.*]] = mhlo.constant dense<-0.000000e+00> : tensor<f32>
  // CHECK-NEXT: %[[red1:.*]] = mhlo.reduce(%[[cmul]] init: %[[init]]) applies mhlo.add across dimensions = [0, 1, 2] : (tensor<8x8x8x8xf32>, tensor<f32>) -> tensor<8xf32>
  // CHECK-NEXT: %[[scr2:.*]] = mhlo.convert %[[red1]] : tensor<8xf32>

                    " a verification pass to ensure invariants are true."}];

  let options = [
    Option<"legalize_chlo_", "legalize-chlo", "bool", /*default=*/"true",
        "Legalizes intermediate chlo ops to hlo">
    ];

  let constructor = "mlir::mhlo::CreateVerifyTFXLALegalizationPass()";
}

def TFXLADeviceSpecificTransforms : Pass<"tfxla-device-specific-transforms",

    %9 = "mhlo.compare"(%arg1, %arg3) {comparison_direction = #mhlo<comparison_direction EQ>} : (tensor<f32>, tensor<f32>) -> tensor<i1>
    %10 = "mhlo.compare"(%arg2, %arg4) {comparison_direction = #mhlo<comparison_direction LT>} : (tensor<i32>, tensor<i32>) -> tensor<i1>
    %11 = mhlo.and %9, %10 : tensor<i1>
    %12 = mhlo.or %7, %11 : tensor<i1>
    %13 = "mhlo.select"(%12, %arg2, %arg4) : (tensor<i1>, tensor<i32>, tensor<i32>) -> tensor<i32>

  "tf.XlaSendToHost"(%arg0) {key = "send2"} : (tensor<i32>) -> ()
  func.return
}

// CHECK: func private @callee([[CALLEE_ARG0:%.*]]: !mhlo.token) -> !mhlo.token
func.func private @callee() {
  // CHECK-NOT:  mhlo.create_token

  // CHECK:      [[ZERO:%.*]] = mhlo.constant dense<0>
  %0 = mhlo.constant dense<0> : tensor<i32>

  // CHECK: %[[QUANT0:.*]] = mhlo.uniform_quantize %[[ARG:.*]] : (tensor<?xf32>) -> tensor<?x!quant.uniform<i8:f32, 1.000000e+00:3>>
  // CHECK: %[[CONVERT_1:.*]] = mhlo.bitcast_convert %[[QUANT0]] : (tensor<?x!quant.uniform<i8:f32, 1.000000e+00:3>>) -> tensor<?xi8>
  // CHECK: mhlo.while()
  // CHECK: cond
  // CHECK: %[[CHECK_RES:.*]] = chlo.broadcast_compare
  // CHECK: mhlo.return %[[CHECK_RES]] : tensor<i1>

#include "tensorflow/compiler/mlir/tf2xla/transforms/passes.h"
#include "xla/mlir_hlo/mhlo/IR/hlo_ops.h"
#include "xla/mlir_hlo/mhlo/IR/register.h"
#include "xla/mlir_hlo/mhlo/transforms/passes.h"
#include "xla/mlir_hlo/mhlo/transforms/rewriters.h"
#include "xla/mlir_hlo/mhlo/utils/type_conversion.h"

namespace mlir {
namespace odml {

class TFToMhloPass

  // CHECK-DAG: %[[RES:.*]] = mhlo.constant dense<{{\[\[\[\[}}1, 4, 9, 16], [5, 12, 21, 32]]]]> : tensor<1x1x2x4xi32>
  %cst0 = mhlo.constant dense<[1, 2, 3, 4]> : tensor<4xi32>
  %cst1 = mhlo.constant dense<[[[[1, 2, 3, 4], [5, 6, 7, 8]]]]> : tensor<1x1x2x4xi32>
  %0 = "mhlo.broadcast_in_dim"(%cst0) <{broadcast_dimensions = dense<3> : tensor<1xi64>}> : (tensor<4xi32>) -> tensor<1x1x2x4xi32>
  %1 = mhlo.multiply %0, %cst1 : tensor<1x1x2x4xi32>

  // CHECK-DAG: %[[FILTER:.+]] = mhlo.constant dense<{{\[\[\[\[}}1.000000e+00, 2.000000e+00], [3.000000e+00, 4.000000e+00], [5.000000e+00, 6.000000e+00]]]]> : tensor<1x1x3x2xf32>
  // CHECK-DAG: %[[CST_0:.+]] = mhlo.constant dense<[1.000000e-01, 2.000000e-01]> : tensor<2xf32>
  // CHECK-DAG: %[[CST_1:.+]] = mhlo.constant dense<[3.000000e-01, 4.000000e-01]> : tensor<2xf32>

    %0 = mhlo.constant dense<42> : tensor<i32>  // padding value
    %1 = mhlo.constant dense<[-10, 11]> : tensor<2xi32>  // k
    // CHECK: mhlo.iota
    // CHECK: mhlo.reshape
    // CHECK: mhlo.concatenate
    // CHECK: mhlo.gather
    // CHECK: mhlo.broadcast
    // CHECK: mhlo.select
    %2 = "tf.MatrixDiagPartV3"(%arg0, %1, %0) {

        return
      }
    ```

    This might be legalized like this:

    ```mlir
      func @send_to_host(%arg0: tensor<i32>) {
        %0 = "mhlo.create_token"() : () -> !mhlo.token
        %1 = "mhlo.send"(%arg0, %0) {...}} : (tensor<i32>, !mhlo.token) -> !mhlo.token
        return
      }
    '''
  }];