func.func @sharding_custom_call_removed(%arg0: tensor<3xf32>) -> tensor<3xf32> {
  %1 = stablehlo.custom_call @Sharding(%arg0) {mhlo.sharding = ""} : (tensor<3xf32>) -> tensor<3xf32>
  return %1 : tensor<3xf32>
}
// CHECK-NOT: custom_call

// -----

// Tests that a custom_call that is not @Sharding is not removed.

// CHECK-LABEL: custom_call_not_removed
func.func @custom_call_not_removed(%arg0: tensor<3xf32>) -> tensor<3xf32> {

      %outputs_9, %control_10 = tf_executor.island wraps "tf.Mul"(%outputs_7, %outputs) {T = "tfdtype$DT_FLOAT", _tpu_replicate = "cluster", device = "", name = "mul"} : (tensor<0xf32>, tensor<f32>) -> tensor<0xf32>
      %outputs_11, %control_12 = tf_executor.island wraps "tf.Placeholder"() {device = "", dtype = "tfdtype$DT_FLOAT", name = "y", shape = "tfshape$dim { }"} : () -> tensor<0xf32>

func.func @main(%arg0: tensor<4xf32>) -> tensor<4xf32> {
  %0 = "tfl.pseudo_const" () {value = dense<1.0> : tensor<4xf32>} : () -> tensor<4xf32> loc("Const")
  %1 = "tfl.squared_difference"(%arg0, %0) {fused_activation_function = "NONE"} : (tensor<4xf32>, tensor<4xf32>) -> tensor<4xf32> loc("squared_difference")
  %2 = "tfl.mul"(%0, %1) {fused_activation_function = "NONE"} : (tensor<4xf32>, tensor<4xf32>) -> tensor<4xf32> loc("mul")

  %tmp0, %ctl0 = tfl.control_node controls "tfl.neg"(%arg0): (tensor<1xf32>) -> tensor<1xf32>
  %tmp1, %ctl1 = tfl.control_node controls "tfl.neg"(%tmp0): (tensor<1xf32>) -> tensor<1xf32>
  %tmp2, %ctl2 = tfl.control_node controls "tfl.neg"(%tmp1): (tensor<1xf32>) -> tensor<1xf32>
  %tmp3, %ctl3 = tfl.control_node controls "tfl.neg"(%tmp2): (tensor<1xf32>) -> tensor<1xf32>

module {
func.func @main(%arg0: tensor<2xi32>) -> tensor<2xi32> attributes {tf.entry_function = {inputs = "arg0", outputs = "tfl.custom1"}} {
  %0 = "tfl.custom"(%arg0, %arg0) {custom_code = "stablehlo.maximum", custom_option = #tfl<const_bytes : "0x00000100002401">} : (tensor<2xi32>, tensor<2xi32>) -> tensor<2xi32>
  func.return %0 : tensor<2xi32>
}
}

// CHECK:       module {

module {
func.func @main(%arg0: tensor<2xi32>) -> tensor<2xi32> attributes {tf.entry_function = {inputs = "arg0", outputs = "tfl.custom1"}} {
  %0 = "tfl.custom"(%arg0, %arg0) {custom_code = "stablehlo.multiply", custom_option = #tfl<const_bytes : "0x00000100002401">} : (tensor<2xi32>, tensor<2xi32>) -> tensor<2xi32>
  func.return %0 : tensor<2xi32>
}
}

// CHECK:       module {

  func.return %1 : tensor<4x4x4xf32>
}

func.func @testMultipleLstms(%arg0: tensor<4x4xf32>, %arg1: tensor<4xf32>, %arg2: tensor<4x4x4xf32>) -> tensor<4x4x4xf32> {

module {
func.func @main(%arg0: tensor<2xi32>) -> tensor<2xi32> attributes {tf.entry_function = {inputs = "arg0", outputs = "tfl.custom1"}} {
  %0 = "tfl.custom"(%arg0, %arg0) {custom_code = "stablehlo.add", custom_option = #tfl<const_bytes : "0x00000100002401">} : (tensor<2xi32>, tensor<2xi32>) -> tensor<2xi32>
  func.return %0 : tensor<2xi32>
}
}

// CHECK:       module {

  return %0 : tensor<1x1600x1x1xf32>
}

//CHECK:func.func private @custom_tf_op(%arg0: tensor<1x320x1x1xf32>, %arg1: tensor<2xi32>) -> tensor<1x1600x1x1xf32> {
//CHECK: %0 = "vhlo.custom_call_v1"(%arg0, %arg1) <{

func.func @cond_true(%arg0: tensor<*xf32>, %arg1: tensor<*xf32>) -> tensor<*xf32> {
  %0 = tf_executor.graph {
    %outputs, %control = tf_executor.island wraps "tf.Add"(%arg0, %arg1) : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32>
    tf_executor.fetch %outputs : tensor<*xf32>
  }
  func.return %0 : tensor<*xf32>
}

func.func @cond_false(%arg0: tensor<*xf32>, %arg1: tensor<*xf32>) -> tensor<*xf32> {
  %0 = tf_executor.graph {