module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 351 : i32}} {

  if (lower_to_xla_hlo) {
    // This pass operates on MHLO control flow ops so it should be legalized
    // after the control flow ops are legalized.
    pm.addPass(mlir::mhlo::CreateLegalizeTFCommunicationPass());

    // Everything should be MHLO after this.
    if (!allow_partial_conversion) {
      pm.addNestedPass<mlir::func::FuncOp>(
          mlir::mhlo::CreateVerifyTFXLALegalizationPass(legalize_chlo));
    }
  }

  ];
  let dependentDialects = [
    "mlir::stablehlo::StablehloDialect", "mlir::TF::TensorFlowDialect",
    "mlir::func::FuncDialect", "mlir::mhlo::MhloDialect",
    "mlir::quant::QuantizationDialect", "mlir::chlo::ChloDialect",
    "mlir::vhlo::VhloDialect", "mlir::shape::ShapeDialect",
    "mlir::quantfork::QuantizationForkDialect",
  ];
}

  // Legalize MHLO to StableHLO should be moved closer to where it is needed
  // There are some entry points that start with HLO->MHLO like
  // jax_to_tfl_flatbuffer.cc which can likely be updated to emit StableHLO
  // to be consistent with other entrypoints.
  pass_manager.addPass(mlir::mhlo::createHloLegalizeToStablehloPass());

  // Decompose CHLO into StableHLO ops

#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
#include "mlir/IR/Location.h"  // from @llvm-project
#include "mlir/IR/Types.h"  // from @llvm-project
#include "xla/mlir_hlo/mhlo/IR/hlo_ops.h"

namespace mlir {
namespace mhlo {

// Builds body for reduce op by using the template binary op as the
// reducer op.
template <typename Op>
void BuildReduceBody(Type element_type, Region* body, OpBuilder* builder) {

// RUN: tf-mlir-translate -mlir-tf-to-hlo-text %s -tf-input-shapes=128,10:10,1024:128,1024 -tf-xla-emit-use-tuple-args -tf-xla-emit-return-tuple | FileCheck %s

module attributes {tf.versions = {producer = 179 : i32}} {
  func.func @main(%arg0: tensor<128x10xf32> {mhlo.sharding = "\08\03\1A\02\01\02\22\02\00\01"}, %arg1: tensor<10x1024xf32> {mhlo.sharding = "\08\01\1A\01\01\22\01\00"}, %arg2: tensor<128x1024xf32> {mhlo.sharding = ""}) {
    func.return
  }
}

include "mhlo/IR/hlo_ops.td"


// See the function doc in the header file.
def GetPadOpType : NativeCodeCall<
  "GetPadOpType((*$0.begin()).getDefiningOp<mhlo::CompositeOp>())">;

// See the function doc in the header file.
def GetAvgPoolOpPadAttr: 
  NativeCodeCall<"GetAvgPoolOpPadAttr($_builder, (*$0.begin()).getDefiningOp<mhlo::CompositeOp>())">;

  auto iota_op = dyn_cast_or_null<mhlo::IotaOp>(iota.getDefiningOp());
  if (!iota_op || dimensions.getNumElements() != 1) return false;
  auto dim = *dimensions.value_begin<APInt>();
  return dim == iota_op.getIotaDimension();
}

// It matches %iota generated from the following mlir codes:
//
// %iota_r1 = mhlo.constant dense<[0, 1, 2, ..., L]>
// %iota = "mhlo.broadcast_in_dim(%iota_r1){

      return %1: tensor<2048xi32>
   }
   // CHECK-LABEL: func @tpu_func
   // CHECK: mhlo.type_extensions
   func.func @tpu_func (
      %arg0: tensor<2048xi32> {mhlo.sharding = "\08\01\1A\01\01\22\01\00"}, %arg1: tensor<2048xi32> {mhlo.sharding = "\08\01\1A\01\01\22\01\00"}) -> (tensor<2048xi32> {mhlo.sharding = "\08\01\1A\01\01\22\01\00"}) {
  // TODO(b/292540052): Below tf.addV2 instruction is replaced with just

#include "tensorflow/compiler/mlir/lite/stablehlo/transforms/hlo_matchers.h"
#include "xla/mlir_hlo/mhlo/IR/hlo_ops.h"

namespace mlir {
namespace odml {

LogicalResult MatchReduceToArgMinMaxType1(mhlo::ReduceOp reduce_op,
                                          bool is_float, bool is_argmax);

LogicalResult MatchReduceToArgMinMaxType2(mhlo::ReduceOp reduce_op,
                                          bool is_argmax);