func.return %0#0 : tensor<8x8x8xf32>
}

// -----

// Test invalid tf.FusedBatchNorm
func.func @testFusedBatchNormWrongScaleType(tensor<8x8x8x8xf32>, tensor<8xi32>, tensor<8xf32>, tensor<8xf32>, tensor<8xf32>) -> tensor<8x8x8x8xf32> {

    %0 = "tf.Sub"(%cst_0, %cst_1) : (tensor<8xf32>, tensor<8xf32>) -> tensor<8xf32>
    return %0 : tensor<8xf32>
  }
// CHECK: @serving_default
// CHECK-DAG: %[[VAR_HANDLE_2:.*]] = "tf.VarHandleOp"() {{.*shared_name = "const_0".*}} : () -> tensor<!tf_type.resource<tensor<8xf32>>>
// CHECK-DAG: %[[READ_VAR_0:.*]] = "tf.ReadVariableOp"(%[[VAR_HANDLE_2]]) : (tensor<!tf_type.resource<tensor<8xf32>>>) -> tensor<8xf32>

  // Use other output
  %3:6 = "tf.FusedBatchNormV3"( %2#0, %arg1, %arg2, %arg3, %arg4) {T = "tfdtype$DT_FLOAT", U = "tfdtype$DT_FLOAT", data_format = "NHWC", epsilon = 0.001 : f32, is_training = false} : (tensor<8x8x8x8xf32>, tensor<8xf32>, tensor<8xf32>, tensor<8xf32>, tensor<8xf32>) -> (tensor<8x8x8x8xf32>, tensor<8xf32>, tensor<8xf32>, tensor<8xf32>, tensor<8xf32>, tensor<8xf32>)

// CHECK:           return %[[VAL_2]], %[[VAL_3]] : tensor<4xf32>, tensor<4xf32>
// CHECK:         }
func.func @broadcast_atan2(%arg0: tensor<1xf32>, %arg1: tensor<4xf32>) -> (tensor<4xf32>, tensor<4xf32>) {
  %0 = "mhlo.broadcast_in_dim"(%arg0) <{broadcast_dimensions = dense<[0]> : tensor<1xi64>}> : (tensor<1xf32>) -> tensor<4xf32>
  %1 = mhlo.atan2 %0, %arg1 : tensor<4xf32>
  %2 = mhlo.atan2 %arg1, %0 : tensor<4xf32>

  func.return %2 : tensor<4xf32>
}

// -----

// CHECK-LABEL: func @uniform_quantize_requantize_and_dequantize_per_axis
func.func @uniform_quantize_requantize_and_dequantize_per_axis(%arg0 : tensor<2x2xf32>) -> tensor<2x2xf32> {
  %scales_0 = "tf.Const"() { value = dense<[1.0, 2.0]> : tensor<2xf32> } : () -> tensor<2xf32>

  %0 = stablehlo.constant dense<2.000000e+00> : tensor<4x2x3x3xf32>  // weight
  %1 = stablehlo.constant dense<3.000000e+00> : tensor<4xf32>  // bias
  %2 = stablehlo.constant dense<0.000000e+00> : tensor<1x4x5x5xf32>  // relu
  %3 = stablehlo.broadcast_in_dim %1, dims = [1] : (tensor<4xf32>) -> tensor<1x4x5x5xf32>

  EXPECT_THAT(num_bytes, Eq(3));
}

TEST_F(GetSizeInBytesTest, Float32ConstOpSizeInBytes) {
  constexpr absl::string_view kConstOpExpr =
      R"mlir(%cst = "tf.Const"() {value = dense<3.0> : tensor<4xf32>} : () -> tensor<4xf32>)mlir";

  Block block{};
  TF::ConstOp int_tensor_const_op = ParseConstOp(kConstOpExpr, block, ctx_);

  const int64_t num_bytes = GetSizeInBytes(int_tensor_const_op);
  EXPECT_THAT(num_bytes, Eq(16));
}

// expected-error@-1 {{'tf_executor.Switch' op operand #1 must be tensor of 1-bit signless integer values}}
    tf_executor.fetch %true : tensor<*xf32>
  }
  func.return %result : tensor<*xf32>
}

// -----

// Check that a switch result type matches the input type.
func.func @invalid_switch(%arg0: tensor<*xf32>, %arg1: tensor<i1>) -> tensor<*xf32> {

func.func @add_with_activation_transpose_broadcasted_rhs(%arg0: tensor<1x3x3x4xf32>) -> tensor<1x4x3x3xf32> {
  %0 = stablehlo.constant dense<2.000000e+00> : tensor<4xf32>
  %1 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<4xf32>) -> tensor<1x4x3x3xf32>
  %2 = stablehlo.transpose %arg0, dims = [0, 3, 1, 2] : (tensor<1x3x3x4xf32>) -> tensor<1x4x3x3xf32>
  %3 = stablehlo.add %2, %1 : tensor<1x4x3x3xf32>

  %0 = "quantfork.stats"(%arg0) {
    layerStats = dense<[-1.28e-5, 1.27e-5]> : tensor<2xf32>
  } : (tensor<1x5x5x2xf32>) -> tensor<1x5x5x2xf32>
  %w = arith.constant dense<[[[[-1.0, 1.0]]], [[[1.0, 2.0]]], [[[-2.0, 1.0]]]]> : tensor<3x1x1x2xf32>
  %b = arith.constant dense<[1.0e-2, 2.1473647e1, -2.1473647e2]> : tensor<3xf32>