%cst = "tf.Const"() <{value = dense<1> : tensor<i32>}> {device = ""} : () -> tensor<i32>

func.func private @conv(%input: tensor<1x3x4x3xf32> {tf._user_specified_name = "input_tensor"}) -> tensor<*xf32> attributes {tf._construction_context = "kEagerRuntime", tf._input_shapes = [#tf_type.shape<1x3x4x3>]} {
  %weight = arith.constant dense_resource<__elided__> : tensor<2x3x3x2xf32>
  %bias = arith.constant dense<[7.11401462, 7.05456924]> : tensor<2xf32>

// Test hoist var handle op and read variable op in the batch function.

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

  func.func private @cond_false_80(%arg0: tensor<1x4xf32> {tf._user_specified_name = "x"}) -> (tensor<i1>, tensor<1x3xf32>) attributes {tf._construction_context = "kEagerRuntime", tf._input_shapes = [#tf_type.shape<1x4>], tf._original_func_name = "cond_false_8"} {
    %cst = "tf.Const"() <{value = dense<true> : tensor<i1>}> {device = ""} : () -> tensor<i1>

    return %3 : tensor<1x3xf32>
  }


  func.func private @cond_false_80(%arg0: tensor<1x4xf32> {tf._user_specified_name = "x"}) -> (tensor<i1>, tensor<1x3xf32>) attributes {tf._construction_context = "kEagerRuntime", tf._input_shapes = [#tf_type.shape<1x4>], tf._original_func_name = "cond_false_8"} {
    %cst = "tf.Const"() <{value = dense<true> : tensor<i1>}> {device = ""} : () -> tensor<i1>

        specs_.unconditionally_use_set_output_shapes;
    for (const auto& name_and_value : func_def->attr()) {
      if (name_and_value.first == "_input_shapes") {
        auto& list = name_and_value.second.list();
        auto& signature = func_def->signature();
        // Some models have "_input_shapes" attribute, but with its value empty
        if (list.shape_size() > 0 &&
            list.shape_size() != signature.input_arg_size()) {

class ShapeInference
    : public impl::TensorFlowShapeInferencePassBase<ShapeInference> {
 public:
  ShapeInference() = default;
  explicit ShapeInference(ArrayRef<ArrayRef<int64_t>> input_shapes)
      : input_shapes_(input_shapes) {}
  void runOnOperation() override {
    // Parse `input_arg_shapes_` if provided (test only)
    SmallVector<ArrayRef<int64_t>> input_shapes_vec;

    // Create input_shapes.
    TF_ShapeAndTypeList* input_shapes =
        TF_NewShapeAndTypeList(input_shapes_vec.size());
    for (size_t i = 0; i < input_shapes_vec.size(); ++i) {
      const auto& input_shape = input_shapes_vec[i];
      if (input_shape.has_value()) {
        TF_ShapeAndTypeListSetShape(input_shapes, i, input_shape->data(),
                                    input_shape->size());

  // Set input_shapes.
  for (int i = 0; i < num_inputs; ++i) {
    std::vector<DimensionHandle> dims;
    const TF_ShapeAndType& input_shape = input_shapes->items[i];
    if (input_shape.num_dims == InferenceContext::kUnknownRank) {
      c.SetInput(i, c.UnknownShape());
      continue;
    }
    dims.reserve(input_shape.num_dims);
    for (int j = 0; j < input_shape.num_dims; ++j) {

        }
        attr {
          key: "_construction_context"
          value {
            s: "kEagerRuntime"
          }
        }
        attr {
          key: "_input_shapes"
          value {
            list {
              shape {
                dim {
                  size: 150
                }
                dim {
                  size: 150
                }