}) {sym_name = "__inference_wrapped_function_4260250_0", sym_visibility = "private", tf._input_shapes = [#tf_type.shape<>], tf.signature.is_stateful, function_type = (tensor<i32>) -> ()} : () -> ()
  "func.func"() ({
  ^bb0(%arg0: tensor<i32>):
    "func.return"(%arg0) : (tensor<i32>) -> ()

    attr {
      key: "_input_shapes"
      value {
      }
    }
  }
}
versions {
  producer: 85
  min_consumer: 12
}

# Match partitioned call in main and capture the callee name.
# CHECK-LABEL: func @main
# CHECK: tf.PartitionedCall
# CHECK-SAME: f = @[[FUNCTION:[a-zA-Z0-9_]*]]

# Verify that callee has the unit attribute tf._input_shapes.
# CHECK: func private @[[FUNCTION]]

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

# Verify that the _input_shapes attribute of the FunctionDef is respected.
# This also checks that the output type is correctly inferred based on
# that.
#CHECK: func private @identity_function0(%arg0: tensor<i32>) -> tensor<i32>

node {
  name: "Placeholder"
  op: "Placeholder"
  attr {
    key: "dtype"
    value {
      type: DT_BOOL

func.func private @whitespace_tokenizer_rank1(%arg0: tensor<1x!tf_type.string> {tf._user_specified_name = "input"}) -> (tensor<?x!tf_type.string>, tensor<?xi64>) attributes {tf._input_shapes = [#tf_type.shape<1>], tf._implements = #tf_type.func<@"tftext:WhitespaceTokenizer", {}>, tf.signature.is_stateful} {
  %0 = "tf.Const"() {value = dense<[0, 1]> : tensor<2xi64>} : () -> tensor<2xi64>

// CHECK-LABEL: func private @batched_function
// CHECK: arg1
func.func private @batched_function(%arg0: tensor<1x3xf32>, %arg1: tensor<*x!tf_type.resource>) -> tensor<1x3xf32>
  attributes {tf._input_shapes = [#tf_type.shape<1x3>, #tf_type.shape<*>], tf.signature.is_stateful} {
  // CHECK: [[handle:%.*]] = "tf.VarHandleOp"()
  // CHECK: "tf.ReadVariableOp"([[handle]])

      name: "func_name"
      input_arg {
        name: "arg0"
        type: DT_BOOL
      }
    }
    ret {
      key: "retval0"
      value: "arg0"
    }
    attr: {
      key: "_input_shapes"
      value: {
      }
    }
  }

    func.return %val : tensor<0xf32>
  }

  // expected-error@+1 {{Number of arguments and 'tf._input_shapes' attribute size do not match. Num args: 2, tf._input_shapes size: 3}}
  func.func private @f_callee(%arg0: tensor<0xf32>, %arg1: tensor<*x!tf_type.resource>) -> tensor<0xf32> attributes {tf._input_shapes = [#tf_type.shape<0>, #tf_type.shape<>, #tf_type.shape<9x9x9>]} {

                             %arg4: tensor<10xf32> {tf._user_specified_name = "b2"}) -> tensor<10xf32>
          attributes {tf._construction_context = "kEagerRuntime", tf._input_shapes = [#tf_type.shape<32x28x28x1>, #tf_type.shape<3x3x1x5>, #tf_type.shape<5>, #tf_type.shape<3920x10>, #tf_type.shape<10>]} {
  return %arg4 : tensor<10xf32>
}

// -----

// CHECK-LABEL: strips_attributes_with_tf_values

func.func private @batched_function(%arg0: tensor<1x3xf32> {tf._user_specified_name = "0"}, %arg1: tensor<*x!tf_type.resource>) -> tensor<1x3xf32> attributes {tf._input_shapes = [#tf_type.shape<1x3>, #tf_type.shape<*>], tf.signature.is_stateful} {
  %0 = "tf.ReadVariableOp"(%arg1) {device = "/device:CPU:0"} : (tensor<*x!tf_type.resource>) -> tensor<1x3xf32>

    }
    func.return %0 : tensor<*xf32>
  }
  func.func private @__inference_call_440(%arg0: tensor<?x28x28x3xf32> {tf._user_specified_name = "inputs"}) -> tensor<*xf32> attributes {tf._input_shapes = [#tf_type.shape<?x28x28x3>], tf.signature.is_stateful} {
    %0 = tf_executor.graph {
      %outputs, %control = tf_executor.island wraps "tf.Const"() {value = dense<0.000000e+00> : tensor<32xf32>} : () -> tensor<32xf32>