%1 = "tf.AddN"(%arg0, %arg1) : (tensor<!tf_type.variant<tensor<*xi32>>>, tensor<!tf_type.variant<tensor<*xi32>>>) -> tensor<!tf_type.variant<tensor<*xi32>>>

      f = @identity} : (tensor<!tf_type.variant<tensor<f32>>>) -> tensor<!tf_type.variant<tensor<f32>>>
  %2 = "tf.OptionalGetValue"(%1) : (tensor<!tf_type.variant<tensor<f32>>>) -> tensor<f32>
  %3 = "tf.OptionalNone"() : () -> tensor<!tf_type.variant<tensor<f32>>>
  return %2, %3 : tensor<f32>, tensor<!tf_type.variant<tensor<f32>>>
}

func.func private @identity(%arg0: tensor<!tf_type.variant<tensor<f32>>>) -> tensor<!tf_type.variant<tensor<f32>>> {

  %outputs_60 =  "tf.Identity"(%outputs_58) {device = ""} : (tensor<*x!tf_type.variant>) -> tensor<*x!tf_type.variant>
  return %outputs_34, %outputs_36, %outputs_44, %outputs_60, %arg4, %arg5 : tensor<*xi32>, tensor<*xi32>, tensor<*xi32>, tensor<*x!tf_type.variant>, tensor<!tf_type.variant>, tensor<?xf32>
}

  func.return
}

// CHECK-LABEL: func @variant
func.func @variant(%arg: tensor<!tf_type.variant>) {
  // CHECK: tf.ZerosLike
  // CHECK-SAME: T = !corert.variant
  %0 = "tf.ZerosLike"(%arg) {device = "/device:CPU:0", T = !tf_type.variant} : (tensor<!tf_type.variant>) -> tensor<!tf_type.variant>
  func.return
}

// Checks that TF quantized attrs are lowered to the corert types

}

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",

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

//===----------------------------------------------------------------------===//
// Matmul op legalizations.

        %3 = "tf.OptionalHasValue"(%2) : (tensor<!tf_type.variant>) -> tensor<i1>
        %4 = "tf.IfRegion"(%3) ({
          %5 = "tf.OptionalGetValue"(%2) : (tensor<!tf_type.variant>) -> tensor<32xf32>

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

// CHECK:...