if (!type || type.getElementType().isUnsignedInteger(64)) {
      return rewriter.notifyMatchFailure(dynamic_iota_op,
                                         "TF::RangeOp doesn't support UI64");
    }
    // Only support 1D for now.
    if (type.getRank() > 1 || dynamic_iota_op.getIotaDimension() != 0) {
      return rewriter.notifyMatchFailure(
          dynamic_iota_op, [&](::mlir::Diagnostic& diag) {

  func.return %2, %5 : tensor<?xi32>, tensor<?xi64>
}

// -----

// CHECK-LABEL: func @two_clusters
func.func @two_clusters(%arg0: tensor<*xi32>, %arg1: tensor<?xi64>) -> (tensor<?xi32>, tensor<?xi64>) {
  // CHECK: %{{[a-z0-9_]*}} = "tf.Const"(){{.*}}value = dense<1>

  // CHECK-SAME: begin_mask = 7 : i64
  // CHECK-SAME: ellipsis_mask = 0 : i64
  // CHECK-SAME: end_mask = 14 : i64
  // CHECK-SAME: new_axis_mask = 0 : i64
  // CHECK-SAME: shrink_axis_mask = 0 : i64

  %0 = "tf.TensorListReserve"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<!tf_type.variant<tensor<*xi64>>>
  // CHECK: %0 = "tfl.custom"(%arg0, %arg1) <{custom_code = "TensorListReserve", custom_option = #tfl<const_bytes : "0x04">}> : (tensor<i32>, tensor<i32>) -> tensor<!tf_type.variant<tensor<*xi64>>>
  func.return %0 : tensor<!tf_type.variant<tensor<*xi64>>>
}

// -----

  // CHECK-NEXT:  return [[FIND]] : tensor<*xi64>
  func.return %1 : tensor<*xi64>
}

// -----

// Test for case with int64 -> string hashtable.
func.func @hashtable_int64_to_string(%arg0: tensor<i64>) -> tensor<*x!tf_type.string> {
  %cst = arith.constant dense<"f"> : tensor<!tf_type.string>

// CHECK-SAME: Tdense = [f32, f32], dense_shapes = [#corert.shape<>, #corert.shape<>]
// CHECK-SAME: num_sparse = 2 : i64, ragged_split_types = [], ragged_value_types = []
// CHECK-SAME: sparse_types = [!corert.string, i64]}
// CHECK-SAME: num_args(7)

// CHECK: tfrt_fallback_async.createop(%0) key(1) device("/device:CPU:0") "tf.ReadVariableOp"() {dtype = f32} num_args(1)

  func.return %0 : tensor<1x1x5x6x7xf32>
  // CHECK-LABEL: strided_slice_big_dims

  %2 = "tf.Shape"(%arg0) : (tensor<?x1x2x?xf32>) -> tensor<4xi32>
  %3 = "tf.StridedSlice"(%2, %1, %0, %1) {begin_mask = 0 : i64, ellipsis_mask = 0 : i64, end_mask = 0 : i64, new_axis_mask = 0 : i64, shrink_axis_mask = 0 : i64} : (tensor<4xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<2xi32>
  func.return %3 : tensor<2xi32>

  %1:4 = "tf.CustomAggregator"(%arg0) {id = "1", calibration_method = 1 : i32, num_bins = 0 : i32, max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32} : (tensor<8x8x8x8xf32>) -> (tensor<8x8x8x8xf32>, tensor<f32>, tensor<f32>, tensor<*xi64>)
  func.return %0#0, %1#0 : tensor<8x8x8x8xf32>, tensor<8x8x8x8xf32>
}
// CHECK: func @customAggregator

    // CHECK-NEXT: %[[UDQ:.*]] = "tf.UniformDequantize"(%[[UQ]], %[[CST_0]], %[[CST_1]]) <{quantization_axis = -1 : i64, quantization_max_val = 127 : i64, quantization_min_val = -128 : i64}> : (tensor<1x!tf_type.qint8>, tensor<f32>, tensor<i32>) -> tensor<1xf32>