// CHECK: %[[DQ_2:.*]] = "tfl.dequantize"({{.*}}) : (tensor<1x1xf16>) -> tensor<1x1xf32>
  // CHECK: %[[DQ_3:.*]] = "tfl.dequantize"({{.*}}) : (tensor<1x1xf16>) -> tensor<1x1xf32>
  // CHECK: %[[DQ_4:.*]] = "tfl.dequantize"({{.*}}) : (tensor<1x1xf16>) -> tensor<1x1xf32>
  // CHECK: %[[DQ_5:.*]] = "tfl.dequantize"({{.*}}) : (tensor<1x1xf16>) -> tensor<1x1xf32>

// CHECK:  return %[[V1]] : tensor<32x1x256xbf16>
}

// CHECK-LABEL: Relu_bf16
func.func @Relu_bf16(%arg0: tensor<2x3xbf16>) -> tensor<2x3xbf16> {
  %cst = arith.constant dense<0.0> : tensor<2x3xbf16>
  %0 = "tfl.maximum"(%arg0, %cst) : (tensor<2x3xbf16>, tensor<2x3xbf16>) -> tensor<2x3xbf16>
  func.return %0 : tensor<2x3xbf16>

  // CHECK: %[[RESULT:.*]] = "tfl.relu"(%arg0)
  // CHECK: return %[[RESULT]]
}

//
module {
  func.func @serving_default() -> tensor<1x3xbf16> {
    %cst = "tf.Const"() <{value = dense<[1]> : tensor<1xi32>}> : () -> tensor<1xi32>
    %0 = "tf.VarHandleOp"() <{container = "", shared_name = "Variable"}> : () -> tensor<!tf_type.resource<tensor<2x3xbf16>>>
    %1 = "tf.ResourceGather"(%0, %cst) <{batch_dims = 0 : i64, validate_indices = true}> : (tensor<!tf_type.resource<tensor<2x3xbf16>>>, tensor<1xi32>) -> tensor<1x3xbf16>

  %0 = "tf.BroadcastTo"(%arg0, %arg1) : (tensor<3xbf16>, tensor<2xi64>) -> tensor<3x3xbf16>
  func.return %0: tensor<3x3xbf16>

// CHECK-LABEL: broadcast_to_bf16
// CHECK:  [[CST:%.*]] = arith.constant dense<1.000000e+00> : tensor<3x3xbf16>
// CHECK:  [[MUL:%.*]] = tfl.mul(%arg0, [[CST]]) <{fused_activation_function = "NONE"}> : (tensor<3xbf16>, tensor<3x3xbf16>) -> tensor<3x3xbf16>

    %1 = "tf.Cast"(%0) <{truncate = false}> : (tensor<3x1xf32>) -> tensor<3x1xbf16>
    %2 = "tf.VarHandleOp"() <{container = "", shared_name = "y"}> : () -> tensor<!tf_type.resource<tensor<3x1xbf16>>>
    "tf.AssignVariableOp"(%2, %1) : (tensor<!tf_type.resource<tensor<3x1xbf16>>>, tensor<3x1xbf16>) -> ()
    return
  }

      return %0 : tensor<1x?x4xbf16>
    }
  })";

  TF_ASSERT_OK(LegalizeModule(kModuleWithNonConstParam));
}

TEST_F(Tf2XlaRewriterTest, ErrorsWithInvalidNumberOfParametersToArgs) {

// CHECK: %[[RESULT:.*]] = "tf.FloorMod"(%arg0, %arg1) : (tensor<192x8xbf16>, tensor<192x8xbf16>) -> tensor<192x8xbf16>
// CHECK: return %[[RESULT]]
// CHECK: }
func.func @convert_floor_mod_float(%arg0: tensor<192x8xbf16>, %arg1: tensor<192x8xbf16>) -> tensor<192x8xbf16> {
  %0 = mhlo.constant dense<0.000000e+00> : tensor<192x8xbf16>
  %1 = mhlo.remainder %arg0, %arg1 : tensor<192x8xbf16>

// -----

// CHECK-LABEL: func @max
func.func @max(%arg0: tensor<4x8xf16>) -> tensor<4x1xf16> {
  // CHECK: %[[CAST:.*]] = mhlo.convert %arg0 : tensor<4x8xf16>
  // CHECK: %[[INITIAL:.*]] = mhlo.constant dense<0xFC00> : tensor<f16>
  // CHECK: %[[REDUCED:.*]] = mhlo.reduce(%[[CAST]] init: %[[INITIAL]]) applies mhlo.maximum across dimensions = [1] : (tensor<4x8xf16>, tensor<f16>) -> tensor<4xf16>