}
  func.func @partitioned_called_func(%arg0: tensor<?xi32>) -> (tensor<?xi32>) {
    func.return %arg0 : tensor<?xi32>
  }
  func.func @stateful_partitioned_call_func(%arg0: tensor<?xi32>) -> (tensor<?xi32>) {
    func.return %arg0 : tensor<?xi32>
  }
  func.func @tpu_partitioned_call_func(%arg0: tensor<?xi32>) -> (tensor<?xi32>) {
    func.return %arg0 : tensor<?xi32>
  }

  %0 = "tf.Bitcast"(%arg0) : (tensor<8x16x32x64xf32>) -> tensor<8x16x32x64xf32>
  // CHECK: return %arg0
  func.return %0: tensor<8x16x32x64xf32>
}

// CHECK-LABEL: testDifferentBitcastType
func.func @testDifferentBitcastType(%arg0: tensor<8x16x32x64xf32>) -> tensor<8x16x32x64xi32> {

func.func @testBitcast_v5(%arg0: tensor<3x2xi16>) -> tensor<2xi32> {
  // expected-error @+1 {{the 0th dim of output tensor is 2. It is not equal to the one in input tensor, which is 3}}
  %0 = "tf.Bitcast"(%arg0) : (tensor<3x2xi16>) -> tensor<2xi32>
  func.return %0 : tensor<2xi32>
}

// -----

// CHECK-LABEL: func @testBitcast_v6
func.func @testBitcast_v6(%arg0: tensor<3x2xi32>) -> tensor<3x2x2xi16> {

//   id0 = arg0
//   if-then-branch: [u0,   arg0, arg0]
//   if-else-branch: [arg0, arg0, arg1]
//     => first result is unknown, second and third is passthrough
//   if results    : [*,    arg0, {arg0, arg1}[
//   ID #2: read (unknown)         -> succ {5, 6)
//   ID #3: read (arg0)            -> succ {5}
//   ID #4: read({arg0,arg1})      -> succ {5,6}
//   ID #5: write(arg0)

^bb0(%arg0: tensor<? x f32>):
  // CHECK: "tfl.log"(%arg0)
  %0 = "tfl.log"(%arg0): (tensor<? x f32>) -> tensor<? x f32>
  func.return %0 : tensor<? x f32>
}

// CHECK-LABEL: testLogStaticShapeInputAndDynamicShapeOutput
func.func @testLogStaticShapeInputAndDynamicShapeOutput(tensor<8 x f32>) -> tensor<? x f32> {
^bb0(%arg0: tensor<8 x f32>):
  // CHECK: "tfl.log"(%arg0)

// CHECK:  %3 = "tf.MaxPool"(%arg0)
// CHECK:  %4 = "tf.MaxPool"(%arg0)
}

func.func @abs(%arg0: tensor<8x16xf32>) -> tensor<8x16xf32> {
  %0 = "tf.Abs"(%arg0) : (tensor<8x16xf32>) -> tensor<8x16xf32>
  func.return %0 : tensor<8x16xf32>

// CHECK-LABEL:abs
// CHECK:  "tfl.abs"(%arg0) : (tensor<8x16xf32>) -> tensor<8x16xf32>
}

func.func @any(%arg0: tensor<2x2xi1>, %arg1: tensor<i32>) -> tensor<i1> {

  func.func @replicated_model_parallel_execute_tpu_cluster_func(%arg0: tensor<?xi32>) -> tensor<?xi32> {
    // CHECK: %[[A_OUTPUT:[0-9]*]] = "tf.A"
    %0 = "tf.A"(%arg0) : (tensor<?xi32>) -> tensor<?xi32>
    // CHECK: %[[REPLICATE:[0-9]*]]:4 = tf_device.replicate
    %1:4 = tf_device.replicate([%0, %arg0, %arg0, %arg0] as %ri_0: tensor<?xi32>) {n = 4 : i32} {

func.func @reshape_float(%arg0: tensor<2x3x4xf32>) -> tensor<6x4xf32> {
  %0 = stablehlo.reshape %arg0 : (tensor<2x3x4xf32>) -> tensor<6x4xf32>
  return %0 : tensor<6x4xf32>
}
// CHECK-LABEL: reshape_float
// CHECK-NOT: tfl.reshape
// CHECK: stablehlo.reshape

// -----

// Tests that a quantized `stablehlo.select` is converted to `tfl.select_v2`.

func.func @select(
    %arg0: tensor<1x3xi1>,

module {
// CHECK-LABEL: func private @some_func
// CHECK-LABEL: func @func_with_call
func.func private @some_func(%arg0: tensor<100xf32>) -> tensor<100xf32> attributes {tf.api_implements = "lstm_b4e9f0e7-ac55-42bc-8ef2-8496419a608c"}
func.func @func_with_call(%arg0: tensor<100xf32>) -> tensor<100xf32> {
  %0 = func.call @some_func(%arg0) : (tensor<100xf32>) -> tensor<100xf32>
  func.return %0 : tensor<100xf32>
  }
}

// -----

}

// Builds XlaRecvAtHost node, and replaces all _Arg nodes with it.
absl::StatusOr<Node*> ReplaceArgNodesWithRecvAtHostNode(
    Graph* g, const string& oc_cluster_name,
    std::vector<DataType>* recv_at_host_dtypes, Node* key_placeholder) {
  // TODO(b/77601805): use out nodes for source node, instead of traversing all
  // nodes.
  std::vector<Node*> arg_nodes = GatherNodesWithType(*g, "_Arg");