}

def BatchFunctionOp : TensorflowMlrt_Op<"batch_function", [Pure]> {
  let summary = "Fallback ExecuteOp specialized for tf.BatchFunction.";

  let description = [{
    This kernel executes a variant tf.BatchFunction kernel that supports having
    the `f` attribute as a bytecode function.

    Example:
      %res = tf_mlrt.batch_function(%input, %captured_input)  {
          device = "/device:CPU:0",

    %0 = stablehlo.convert %arg0 : (tensor<1x3x3x4xi8>) -> tensor<1x3x3x4xf32>
    return %0 : tensor<1x3x3x4xf32>
  }
// CHECK: @uniform_dequantize_0
}

// -----

// Tests a variant where there is no stablehlo.convert op in between the
// filter constant and the convolution op.
//
// `filter (f32) -> convolution`
//
// instead of:
//
// `filter (i8) -> convert (i8 -> f32) -> convolution`

    %2 = "tf.Identity"(%1) : (tensor<!tf_type.variant<tensor<f32>>>) -> (tensor<!tf_type.variant<tensor<f32>>>)
    "tf.Identity"(%2) : (tensor<!tf_type.variant<tensor<f32>>>) -> (tensor<!tf_type.variant<tensor<f32>>>)
    tf_device.return
  }) {cluster_attr = "cluster_attr"} : () -> ()
  func.return
}

  func.func @unsupported_dtype(%arg0: tensor<2x!tf_type.variant>) -> tensor<2x!tf_type.variant> {
    // CHECK: tf.AddN
    // expected-remark@+1 {{skipping legalization due to unsupported type 'tensor<2x!tf_type.variant>'}}
    %0 = "tf.AddN"(%arg0, %arg0) : (tensor<2x!tf_type.variant>, tensor<2x!tf_type.variant>) -> tensor<2x!tf_type.variant>

    func.return %0 : tensor<2x!tf_type.variant>
  }

  // CHECK-LABEL: multiple_dialect_ops

  return
}

// -----

func.func @init(%arg0: tensor<4xf32>) -> tensor<7xf32> {
  %0 = builtin.unrealized_conversion_cast to tensor<7xf32>
  return %0 : tensor<7xf32>
}

def BatchFunctionOp : FallbackAsync_Op<"batch_function",
    [Pure, CoreRT_TypedAttributeTrait]> {
  let summary = "Fallback ExecuteOp specialized for tf.BatchFunction.";
  let description = [{
    This kernel executes a variant tf.BatchFunction kernel that supports having
    the `f` attribute as a BEF function.

    Example:
      %res = tfrt_fallback_async.batch_function device("/CPU:0") @f(%input, %captured_input)  {

// CHECK-LABEL: func @testTensorListElementShape
func.func @testTensorListElementShape(%arg0: tensor<!tf_type.variant<tensor<2x4xf32>>>) -> tensor<2xi32> {
  %0 = "tf.TensorListElementShape"(%arg0) : (tensor<!tf_type.variant<tensor<2x4xf32>>>) -> tensor<2xi32>
  // CHECK:         [[cst:%.*]] = "tf.Const{{.*}} dense<{{\[}}2, 4]> : tensor<2xi32>
  // CHECK-NEXT:    return [[cst]] : tensor<2xi32>

  %5 = "tf.Identity"(%4) {__op_key = 10: i32, device = "/device:CPU:0"} : (tensor<!tf_type.variant>) -> tensor<!tf_type.variant>
  func.return %5 : tensor<!tf_type.variant>
}
// CHECK-LABEL: __inference_Dataset_map_lambda_160

    %9 = "tf.opK"(%input_0, %input_4, %9) :
      (tensor<i32>, tensor<!tf_type.variant>, tensor<i1>) -> tensor<i32>
    tf_device.return %9 : tensor<i32>
  }) {device = "DEVICE_ALIAS_0"} : () -> tensor<i32>
  %l = "tf_device.launch"() ( {
    %10 = "tf.opL"(%input_1, %input_4, %9) :
      (tensor<f32>, tensor<!tf_type.variant>, tensor<i1>) -> tensor<f32>
    tf_device.return %10 : tensor<f32>

  INT16 = 7,
  COMPLEX64 = 8,
  INT8 = 9,
  FLOAT64 = 10,
  COMPLEX128 = 11,
  UINT64 = 12,
  // Experimental: Resource and variant types are experimental, that are subject
  // to change. Do not implement custom kernels using resource & variant types
  // now.
  RESOURCE = 13,
  VARIANT = 14,
  UINT32 = 15,
  UINT16 = 16,
  INT4 = 17,
  BFLOAT16 = 18,
}