return %updated_list : tensor<!tf_type.variant<tensor<64x1xbf16>>>
    }
  })";

  auto compilation_result = CompileMlirModule(
      kTensorListSetItemDimensionTooLarge,
      ConfigProto::Experimental::MLIR_BRIDGE_ROLLOUT_UNSPECIFIED);

  // Ensure that it compile

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

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

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

  TF_ASSERT_OK(LegalizeModule(kModuleWithNonConstParam));
}

TEST_F(Tf2XlaRewriterTest, ErrorsWithInvalidNumberOfParametersToArgs) {

  %1 = "tf.MatMul"(%0, %cst) {device = "", transpose_a = false, transpose_b = false} : (tensor<1x10xbf16>, tensor<10x2xbf16>) -> tensor<1x2xbf16>
  %2 = "tf.Cast"(%1) {Truncate = false} : (tensor<1x2xbf16>) -> tensor<1x2xf32>
  %3 = "tf.IdentityN"(%2) {device = ""} : (tensor<1x2xf32>) -> tensor<1x2xf32>
  return %3 : tensor<1x2xf32>
}

// CHECK: func @cast_bf16_matmul_to_fp32

  // CHECK: tfrt_fallback_async.const_dense_tensor  dense<1.000000e+00> : tensor<1xbf16>
  %1 = "tf.Const"() {device = "/device:CPU:0", value = dense<[1.0]> : tensor<1xbf16>} : () -> tensor<4xbf16>
  // CHECK: corert.executeop({{.*}}) "tf.Const"() {dtype = ui64, value = dense<[1, 2, 3, 4]> : tensor<4xui64>} : 1
  %2 = "tf.Const"() {device = "/device:GPU:0", value = dense<[1, 2, 3, 4]> : tensor<4xui64>} : () -> tensor<4xui64>

}

func.func @matmul_batchv3_unknown_dim_bf16(%arg0: tensor<?x4x5xbf16>, %arg1: tensor<5x6xf32>) -> tensor<?x4x6xbf16> {
  %0 = "tf.Cast"(%arg0) : (tensor<?x4x5xbf16>) -> tensor<?x4x5xf32>
  %1 = "tf.BatchMatMulV3"(%0, %arg1) {Ta = "tfdtype$DT_FLOAT", Tb = "tfdtype$DT_FLOAT", device = "/device:CPU:0", name = "MatMul", adj_x = false, adj_y = false} :

func.func @SparseSoftmaxCrossEntropyWithLogits_with_bf16_i64(%features: tensor<2x3xbf16>, %labels: tensor<2xi64>) -> (tensor<2xbf16>, tensor<2x3xbf16>) {
  // CHECK-NOT: tf.SparseSoftmaxCrossEntropyWithLogits
  %0:2 = "tf.SparseSoftmaxCrossEntropyWithLogits"(%features, %labels) : (tensor<2x3xbf16>, tensor<2xi64>) -> (tensor<2xbf16>, tensor<2x3xbf16>)
  func.return %0#0, %0#1 : tensor<2xbf16>, tensor<2x3xbf16>
}

func.func @testValidMatrixBandPartOpCompatibleDynamicShapes(%arg0: tensor<?x10x?xbf16>, %arg1: tensor<i64>, %arg2: tensor<i64>) -> tensor<?x?x8xbf16> {
  %0 = "tf.MatrixBandPart"(%arg0, %arg1, %arg2) : (tensor<?x10x?xbf16>, tensor<i64>, tensor<i64>) -> tensor<?x?x8xbf16>
  func.return %0 : tensor<?x?x8xbf16>
}

// -----

// Test invalid tf.MatrixBandPart