%8 = "tfl.mul"(%2, %3) {fused_activation_function = "NONE"} : (tensor<4xbf16>, tensor<4xbf16>) -> tensor<4xbf16>

  func.return %5, %6, %7, %8 : tensor<bf16>, tensor<4xbf16>, tensor<4xbf16>, tensor<4xbf16>
}

// CHECK-LABEL: @mul_f16
func.func @mul_f16() -> (tensor<f16>, tensor<4xf16>, tensor<4xf16>, tensor<4xf16>) {
  %0 = arith.constant dense<4.5> : tensor<f16>

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

    // FusedMatMul kernel supports limited set of data types.
    Type element_ty = getElementTypeOrSelf(matmul.getType());
    if (!element_ty.isF32() && !element_ty.isBF16()) {
      (void)rewriter.notifyMatchFailure(matmul, [&](Diagnostic &diag) {
        diag << "supported data types for _FusedMatMul are float and bfloat16, "
             << " but got " << element_ty;
      });

// Float16-DAG: %[[b:.*]] = arith.constant dense<0.000000e+00> : tensor<16xf16>
// Float16-DAG: %[[const:.*]] = "tfl.no_value"() <{value}> : () -> none
// Float16-DAG: %[[dq_w:.*]] = "tfl.dequantize"(%[[w]]) : (tensor<3x3x3x8x16xf16>) -> tensor<3x3x3x8x16xf32>
// Float16-DAG: %[[dq_b:.*]] = "tfl.dequantize"(%[[b]]) : (tensor<16xf16>) -> tensor<16xf32>