%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>

}

func.func private @testIf3Then(tensor<*xf32>) -> (tensor<*xf32>, tensor<*xi8>, tensor<*xbf16>)
func.func private @testIf3Else(tensor<*xf32>) -> (tensor<*xf32>, tensor<*xi8>, tensor<*xbf16>)

// CHECK-LABEL: func @testIf3Result(%arg0: tensor<i1>, %arg1: tensor<*xf32>)
func.func @testIf3Result(tensor<i1>, tensor<*xf32>) -> (tensor<*xf32>, tensor<*xi8>, tensor<*xbf16>) {
^bb0(%arg0: tensor<i1>, %arg1: tensor<*xf32>):
  %1:3 = "tf.If"(%arg0, %arg1) {

func.func @single_op_island_forward_block_arg(%arg0: tensor<?x?x?x?xbf16>) -> (tensor<2048xf32>, tensor<?x?x?x?xbf16>) {
  %0:2 = tf_executor.graph {
    %outputs:2, %control = tf_executor.island {
      %1 = "tf.Const"() {value = dense<0.000000e+00> : tensor<2048xf32>} : () -> tensor<2048xf32>
      tf_executor.yield %1, %arg0 : tensor<2048xf32>, tensor<?x?x?x?xbf16>
    }

  %0:2 = tf_executor.graph {
    %outputs, %control = tf_executor.island wraps "tf.Const"() {value = dense<0.000000e+00> : tensor<2048xf32>} : () -> tensor<2048xf32>
    tf_executor.fetch %outputs, %arg0 : tensor<2048xf32>, tensor<?x?x?x?xbf16>
  }
  return %0#0, %0#1 : tensor<2048xf32>, tensor<?x?x?x?xbf16>
}

func.func @single_op_island_forward_block_arg(%arg0: tensor<?x?x?x?xbf16>) -> (tensor<2048xf32>, tensor<?x?x?x?xbf16>) {
  %0:2 = tf_executor.graph {
    %outputs:2, %control = tf_executor.island {
      %1 = "tf.Const"() {value = dense<0.000000e+00> : tensor<2048xf32>} : () -> tensor<2048xf32>
      tf_executor.yield %1, %arg0 : tensor<2048xf32>, tensor<?x?x?x?xbf16>
    }

    %2 = "tf.Cast"(%1) {Truncate = false} : (tensor<?x?x?x?xf32>) -> tensor<?x?x?x?xbf16>
    return %2 : tensor<?x?x?x?xbf16>
  }

// CHECK-LABEL: func @quantize_sharded_weights_with_xladot
// CHECK-DAG: %[[W:.*]] = "tf.Const"() <{value = dense<127> : tensor<512x512xi8>

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

  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>