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

  // CHECK: %[[X_CENTER:.+]] = mhlo.subtract %[[X]], %[[MEAN_BCAST]] : tensor<4x256xf32>
  // CHECK: %[[X_SCALED:.+]] = mhlo.multiply %[[X_CENTER]], %[[SCALE_BCAST]] : tensor<4x256xf32>
  // CHECK: %[[X_NORMED:.+]] = mhlo.divide %[[X_SCALED]], %[[STDDEV_BCAST]] : tensor<4x256xf32>
  // CHECK: %[[RESULT:.+]] = mhlo.add %[[X_NORMED]], %[[OFFSET_BCAST]] : tensor<4x256xf32>
  %0 = "mhlo.batch_norm_inference"(%x, %scale, %offset, %mean, %variance)

  // CHECK-DAG: %[[LHS_1:.*]] = mhlo.convert %arg0 : tensor<3x2xi32>
  // CHECK-DAG: %[[LHS_2:.*]] = mhlo.bitcast_convert %[[LHS_1]] : (tensor<3x2xi32>) -> tensor<3x2x!quant.uniform<i32:f32, 2.000000e+00:4>>
  // CHECK-DAG: %[[RHS_1:.*]] = mhlo.convert %arg1 : tensor<2xi32>
  // CHECK-DAG: %[[RHS_2:.*]] = mhlo.bitcast_convert %[[RHS_1]] : (tensor<2xi32>) -> tensor<2x!quant.uniform<i32:f32, 2.000000e+00:4>>

#include "xla/mlir_hlo/mhlo/IR/hlo_ops.h"  // IWYU pragma: keep

namespace mlir {
namespace odml {

// Given a Composite op that wraps a core.aten.avg_pool2d, returns the padding
// configuration required for the `tfl.pad` if the padding part of the op is
// to be done before average pooling.
DenseIntElementsAttr GetPadOpAttr(Builder& builder, mhlo::CompositeOp op);

#include "tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model.h"
#include "tensorflow/compiler/mlir/tensorflow/transforms/passes.h"
#include "xla/mlir_hlo/mhlo/IR/register.h"
#include "xla/mlir_hlo/mhlo/transforms/passes.h"
#include "tensorflow/core/ir/types/dialect.h"

int main(int argc, char** argv) {
  tensorflow::InitMlir y(&argc, &argv);

  mlir::registerAllPasses();

  // of dynamic shapes. Therefore we lower chlo ops after optimization.
  pm.addNestedPass<func::FuncOp>(CreateOptimizeIntGraphPass());
  pm.addNestedPass<func::FuncOp>(mhlo::createChloLegalizeToHloPass());
  pm.addNestedPass<func::FuncOp>(createCanonicalizerPass());
  pm.addPass(createSymbolDCEPass());
  // MHLO -> StableHLO legalization.

  mlir::tf_saved_model::registerTensorFlowSavedModelPasses();
  mlir::TFL::registerTensorFlowLitePasses();
  mlir::mhlo::registerAllMhloPasses();

  // These are in compiler/mlir/tf2xla and not part of the above MHLO passes.
  mlir::mhlo::registerLegalizeTfPasses();
  mlir::mhlo::registerTfXlaPasses();
  mlir::quant::stablehlo::registerBridgePasses();
  tensorflow::tf2xla::internal::registerTFXLABridgeClusteringPasses();

#include "mlir/Support/LLVM.h"  // from @llvm-project
#include "xla/mlir_hlo/mhlo/IR/hlo_ops.h"

namespace mlir {
namespace odml {
namespace {

#define DEBUG_TYPE "unfold-splat-constant-pass"

#define GEN_PASS_DEF_UNFOLDSPLATCONSTANTPASS
#include "tensorflow/compiler/mlir/lite/stablehlo/transforms/passes.h.inc"

// Undo the MHLO::BroadcastInDimOp folding pattern on splat tensor.

  // on StableHLO, treating StableHLO purely as an input format, and do all
  // optimizations via MHLO passes that can be shared with the OpenXLA compiler.
  // Therefore, this function inserts a StableHLO <=> MHLO roundtrip to make
  // this happen.

  // StableHLO -> MHLO legalization.
  pm.addPass(mhlo::createStablehloLegalizeToHloPass());

  AddMhloOptimizationPasses(pm, /*enable_stablehlo_quantizer=*/false);

}

def LegalizeStablehloToVhloPass : Pass<"stablehlo-legalize-vhlo", "ModuleOp"> {
  let summary = "Legalize StableHLO to VHLO for ODML.";
  let dependentDialects = ["mlir::vhlo::VhloDialect"];
}

def LegalizeVhloToStablehloPass : Pass<"vhlo-legalize-stablehlo", "ModuleOp"> {
  let summary = "Legalize VHLO to StableHLO for ODML.";
  let dependentDialects = ["mlir::stablehlo::StablehloDialect"];
}