}
    Value new_const_op =
        rewriter.create<mhlo::ConstantOp>(bcast_op.getLoc(), result);
    rewriter.replaceOp(bcast_op, {new_const_op});
    return success();
  }
};

using FoldBroadcastInDimBeforeMulOp =
    FoldBroadcastInDimBeforeBinaryElementwiseOp<mhlo::MulOp>;

// Constant folds mhlo.mul, this folder doesn't have an upper limit on how many
// elements can be folded.

  ConversionTarget target(context);
  target.addLegalDialect<TFL::TensorFlowLiteDialect, mhlo::MhloDialect>();
  target.addLegalOp<func::CallOp, func::ConstantOp, arith::ConstantOp>();
  target.addDynamicallyLegalOp<mhlo::CustomCallOp>(IsCustomCallLegal);
  target.addDynamicallyLegalOp<mhlo::ReduceOp>(IsReduceOpLegal);
  // Converted MHLO ops should be marked illegal here.
  // TODO: b/304003568 - Add TF_TransposeOp folding logic to tflite.

  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 "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){

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

// Legalizes TF/XLA communication ops (TF dialect) to HLO dialect communication
// ops.
std::unique_ptr<OperationPass<ModuleOp>> CreateLegalizeTFCommunicationPass();

// Legalizes TF/XLA collective ops (TF dialect) to HLO dialect collective
// ops.
std::unique_ptr<OperationPass<ModuleOp>> CreateLegalizeTFCollectivePass();

// Verifies that the TF/XLA ops have all been lowered to MHLO.

#include "tensorflow/compiler/mlir/tf2xla/transforms/passes.h"
#include "tensorflow/compiler/mlir/tf2xla/transforms/xla_legalize_targets.h"
#include "xla/mlir_hlo/mhlo/IR/hlo_ops.h"
#include "tensorflow/core/lib/monitoring/counter.h"
#include "tensorflow/core/platform/errors.h"

namespace mlir {
namespace mhlo {

namespace {

#define GEN_PASS_DEF_VERIFYTFXLALEGALIZATION
#include "tensorflow/compiler/mlir/tf2xla/transforms/xla_legalize_tf_passes.h.inc"