// In the newly cloned function, check that we have a _tf.If operation and capture the then and else branch.
// CHECK: func private @[[FUNCTIONALIZE_FUNC]]
// CHECK: "tf.If"
// CHECK-SAME:  else_branch = @[[ELSE_FUNC:[A-Za-z0-9_]*]]
// CHECK-SAME:  then_branch = @[[THEN_FUNC:[A-Za-z0-9_]*]]

// We expect the _tf.Add in the else func and the _tf.Mul in the then func

// CHECK: func private @[[ELSE_FUNC]]
// CHECK: "tf.Add"

    }
  }
  attr {
    key: "value"
    value {
      tensor {
        dtype: DT_INT32
        tensor_shape {
        }
        int_val: 0
      }
    }
  }
}
node {
  name: "_tf.foo"
  op: "foo"
  input: "Constant"
}
library {
  function {
    signature {
      name: "foo"
      input_arg {
        name: "arg"
        type: DT_INT32
      }
      output_arg {

    }
  }
  attr {
    key: "value"
    value {
      tensor {
        dtype: DT_INT32
        tensor_shape {
        }
        int_val: 0
      }
    }
  }
}
node {
  name: "_tf.foo"
  op: "foo"
  input: "Constant"
}
library {
  function {
    signature {
      name: "foo"
      input_arg {
        name: "arg"
        type: DT_INT32
      }
      output_arg {

# RUN: tf-mlir-translate -graphdef-to-mlir -tf-enable-shape-inference-on-import=false %s -o - | FileCheck %s

node {
  name: "_tf.PartitionedCall"
  op: "PartitionedCall"
  attr {
    key: "Tin"
    value {
      list {
      }
    }
  }
  attr {
    key: "Tout"
    value {
      list {
        type: DT_FLOAT
      }
    }
  }
  attr {
    key: "f"
    value {
      func {
        name: "foo"
      }

module  {
  tfg.graph #tf_type.version<producer = 62, min_consumer = 12> {
    %Const, %ctl = Const name("Constant") {dtype = i32, value = dense<0> : tensor<i32>} : () -> (tensor<i32>)
    %foo, %ctl_0 = foo(%Const) name("_tf.foo") : (tensor<i32>) -> (tensor<*xi32>)
  }
  tfg.func @foo(%arg: tensor<*xi32> {tfg.name = "arg"})
       -> (tensor<*xi32> {tfg.dtype = i32, tfg.name = "return_value"})
   {
    // CHECK-NOT: ^bb0{{.*}}, %arg2: !tf_type.control

// Creates a pass to convert MLIR to Graph, run user-specified Graph
// Optimization Passes and convert back to MLIR.
// Constraints: This pass expects that all operations in the MLIR module either
// belong to 'tf' or '_tf' dialect. The output is in '_tf' dialect.
class GraphOptPass
    : public mlir::PassWrapper<GraphOptPass,
                               mlir::OperationPass<mlir::ModuleOp>> {