func.func private @add_scatter(%arg0: tensor<*xbf16>, %arg1: tensor<*xbf16>) -> tensor<*xbf16> {
    %0 = "tf.AddV2"(%arg0, %arg1) {device = ""} : (tensor<*xbf16>, tensor<*xbf16>) -> tensor<*xbf16>
    func.return %0 : tensor<*xbf16>
  }

  func.func private @ge_select(%arg0: tensor<*xbf16>, %arg1: tensor<*xbf16>) -> tensor<*xi1> {
    %0 = "tf.GreaterEqual"(%arg0, %arg1) {device = ""} : (tensor<*xbf16>, tensor<*xbf16>) -> tensor<*xi1>

  func.return %0 : tensor<*xbf16>
}

func.func private @sum1(%arg0: tensor<*xbf16>, %arg1: tensor<*xbf16>) -> tensor<*xbf16> {
  %0 = "tf.AddV2"(%arg0, %arg1) : (tensor<*xbf16>, tensor<*xbf16>) -> tensor<*xbf16>
  func.return %0 : tensor<*xbf16>
}

// CHECK-LABEL: testXlaVariadicReduceToV2

  %0 = "tf.Const"() { value = dense<5> : tensor<i32> } : () -> tensor<i32>

  // CHECK: [[VAL:%.+]] = "tf.Const"() <{value =  dense<0.000000e+00> : tensor<5xbf16>}>
  // CHECK: return [[VAL]]
  %1 = "tf.Empty"(%0) : (tensor<i32>) -> (tensor<5xbf16>)
  func.return %1 : tensor<5xbf16>
}

// CHECK-LABEL: func @testShapeN
func.func @testShapeN(%arg0: tensor<f32>, %arg1: tensor<1x32x32x16xf32>) -> (tensor<0xi64>, tensor<4xi64>) {

  if (!tensorType) return false;
  return tensorType.getElementType().isF32();
}

// Returns true iff the given value is a bf16 tensor.
inline bool TFTypeIsBFloat16Tensor(Value value) {
  auto tensorType = mlir::dyn_cast<TensorType>(value.getType());
  if (!tensorType) return false;
  return tensorType.getElementType().isBF16();
}

// Returns true iff the given value is a f16 tensor.
inline bool TFTypeIsHalfTensor(Value value) {

  auto type = mlir::dyn_cast<ShapedType>(val.getType());
  if (!type) return false;
  // Supported original tensor data types.
  if (type.getElementType().isF32() || type.getElementType().isBF16())
    return true;
  return false;
}

std::optional<tensorflow::quantization::QuantizationComponentSpec>
GetWeightComponentSpec(
    const tensorflow::quantization::QuantizationOptions& quantization_options) {

    return absl::OkStatus();
  } else if (type.isF32()) {
    *dtype = DT_FLOAT;
    return absl::OkStatus();
  } else if (type.isF64()) {
    *dtype = DT_DOUBLE;
    return absl::OkStatus();
  } else if (type.isBF16()) {
    *dtype = DT_BFLOAT16;
    return absl::OkStatus();
  } else if (type.isFloat8E4M3FN()) {
    *dtype = DT_FLOAT8_E4M3FN;
    return absl::OkStatus();
  } else if (type.isFloat8E5M2()) {

using absl::StatusOr;

namespace errors = tensorflow::errors;

tflite::TensorType ConvertTypeToTensorType(mlir::Type type) {
  if (type.isF16()) {
    return tflite::TensorType_FLOAT16;
  } else if (type.isBF16()) {
    return tflite::TensorType_BFLOAT16;
  } else if (type.isF32()) {
    return tflite::TensorType_FLOAT32;
  } else if (type.isF64()) {
    return tflite::TensorType_FLOAT64;

}

func.func private @add_scatter(%arg0: tensor<bf16>, %arg1: tensor<bf16>) -> tensor<bf16> {
  %0 = "tf.AddV2"(%arg0, %arg1) {device = ""} : (tensor<bf16>, tensor<bf16>) -> tensor<bf16>
  func.return %0 : tensor<bf16>
}

func.func private @ge_select(%arg0: tensor<bf16>, %arg1: tensor<bf16>) -> tensor<i1> {
  %0 = "tf.GreaterEqual"(%arg0, %arg1) {device = ""} : (tensor<bf16>, tensor<bf16>) -> tensor<i1>

def IsF32ElementType : Constraint<
  CPred<"getElementTypeOrSelf($0).isF32()">>;

// Checks if the value has the type of bfloat16.
def IsBF16ElementType : Constraint<
  CPred<"getElementTypeOrSelf($0).isBF16()">>;

// Checks if the value has the type of UniformQuantizedType.
def IsUniformQuantizedType : Constraint<
  CPred<"getElementTypeOrSelf($0).isa<mlir::quant::UniformQuantizedType>()">>;

  Type type = val_bcast.getType();
  Type elem_type = getElementTypeOrSelf(type);
  // Xla's all_reduce legalizer bitcasts to 32 bits, so only
  // element types size <= 4 bytes are supported.
  if (elem_type.isBF16() || elem_type.isF16() || elem_type.isTF32() ||
      elem_type.isF32()) {
    zero = builder.getFloatAttr(elem_type, 0);
  } else {
    return false;
  }
  if (auto ranked_type = dyn_cast<RankedTensorType>(type)) {