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

  // CHECK: %[[DQ_10:.*]] = "tfl.dequantize"({{.*}}) : (tensor<3xf16>) -> tensor<3xf32>
  // CHECK: %[[DQ_11:.*]] = "tfl.dequantize"({{.*}}) : (tensor<1x3xf16>) -> tensor<1x3xf32>
  // CHECK: %[[DQ_12:.*]] = "tfl.dequantize"({{.*}}) : (tensor<1x3xf16>) -> tensor<1x3xf32>
  // CHECK: %[[lstm:.*]] = "tfl.unidirectional_sequence_lstm"(
  // CHECK-SAME: %arg0,
  // CHECK-SAME: %[[DQ_1]], %[[DQ_2]], %[[DQ_3]], %[[DQ_4]],

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

  TF_ASSERT_OK(LegalizeModule(kModuleWithNonConstParam));
}

TEST_F(Tf2XlaRewriterTest, ErrorsWithInvalidNumberOfParametersToArgs) {

// CHECK: dtype: DT_HALF
// CHECK: half_val: 15360
// CHECK-NEXT: half_val: 16384
    %2:2 = tf_executor.island wraps "tf.Const"() {device = "", dtype = bf16, value = dense<[4.900000e+01, 8.200000e+02]> : tensor<2xbf16>} : () -> tensor<bf16> loc("const3")
// CHECK: name: "const3"
// CHECK-NEXT: op: "Const"
// CHECK: dtype: DT_BFLOAT16
// CHECK: half_val: 16964
// CHECK: half_val: 17485

func.func @floordiv_f16_broadcast(%arg0: tensor<2x3xf16>, %arg1: tensor<3xf16>) -> tensor<2x3xf16> {
  // CHECK-NEXT:  chlo.broadcast_divide
  // CHECK-NEXT:  mhlo.floor
  // CHECK-NEXT:  return
  %0 = "tf.FloorDiv"(%arg0, %arg1) : (tensor<2x3xf16>, tensor<3xf16>) -> tensor<2x3xf16>
  func.return %0: tensor<2x3xf16>
}

// -----

// CHECK-LABEL: func @floordiv_dynamic

// CHECK:           return %[[VAL_4]] : tensor<2x3xf16>
// CHECK:         }
func.func @floordiv_f16_broadcast(%arg0: tensor<2x3xf16>, %arg1: tensor<3xf16>) -> tensor<2x3xf16> {
  %0 = "chlo.broadcast_divide"(%arg0, %arg1) {broadcast_dimensions = array<i64: 1>} : (tensor<2x3xf16>, tensor<3xf16>) -> tensor<2x3xf16>

// -----

func.func @testInvalidSelect(%arg0: tensor<3xi1>, %arg1: tensor<2x3xf16>, %arg2: tensor<2x3xf16>) -> tensor<2x3xf16> {
  // expected-error @+1 {{requires that, when pred is a vector, the shape matches the first dimension of t and e}}
  %0 = "tf.Select"(%arg0, %arg1, %arg2) : (tensor<3xi1>, tensor<2x3xf16>, tensor<2x3xf16>) -> tensor<2x3xf16>
  func.return %0: tensor<2x3xf16>
}

// -----

  }

  // CHECK-LABEL: func @call_in_graph_1
  func.func @call_in_graph_1(%arg0: tensor<?x?x?x?xbf16>, %arg1: tensor<5x5x1x32xbf16>) -> tensor<*xbf16> {
    // CHECK: tf_executor.fetch %outputs : tensor<?x?x?x32xbf16>
    %0 = tf_executor.graph {
      %1:2 = tf_executor.island wraps "tf.PartitionedCall"(%arg0, %arg1) {