auto context = CreateContext();

  EXPECT_EQ(ToBfloat16Type(Float8E4M3FNType::get(context.get())),
            Float8E4M3FNType::get(context.get()));
  EXPECT_EQ(ToBfloat16Type(Float16Type::get(context.get())),
            Float16Type::get(context.get()));
  EXPECT_EQ(ToBfloat16Type(BFloat16Type::get(context.get())),
            BFloat16Type::get(context.get()));
  EXPECT_EQ(ToBfloat16Type(Float32Type::get(context.get())),

            op->getLoc(), ToBfloat16Type(op_operand.get().getType()),
            op_operand.get()));
      }
    }
    builder.setInsertionPointAfter(op);
    for (auto op_result : op->getOpResults()) {
      if (IsLargeFloatType(op_result.getType())) {
        const Type original_type = op_result.getType();
        op_result.setType(ToBfloat16Type(original_type));
        const Value cast =

namespace mlir::quant::stablehlo {

bool IsLargeFloatType(Type type) {
  type = getElementTypeOrSelf(type);
  return isa<FloatType>(type) && type.getIntOrFloatBitWidth() > 16;
}

Type ToBfloat16Type(Type type) {
  if (auto shaped = mlir::dyn_cast<ShapedType>(type)) {
    const Type elem = shaped.getElementType();
    if (IsLargeFloatType(elem)) {
      return shaped.clone(BFloat16Type::get(type.getContext()));
    }

// Returns true if the type or its element type is a float type with bit_width
// > 16.
bool IsLargeFloatType(Type type);

// Converts large float type to bfloat16. Otherwise returns original type.
Type ToBfloat16Type(Type type);

}  // namespace mlir::quant::stablehlo

namespace mlir::quant::stablehlo {
namespace {

class BFloat16TypeConverter : public TypeConverter {
 public:
  BFloat16TypeConverter() {
    addConversion([](const Type type) -> Type {
      return IsLargeFloatType(type) ? ToBfloat16Type(type) : type;
    });
  }
};

// This helper function makes legality check easier. Both convert ops in the
// patterns below are considered legal:
//  - `BitcastConvertOp` (i32 -> f32) + `ConvertOp` (f32 -> bf16)