func.func @f(%arg0: tensor<!tf_type.resource<tensor<f32>>> {tf_saved_model.bound_input = @v}) -> (tensor<f32> {tf_saved_model.index_path = []})
  attributes {tf_saved_model.exported_names = ["f"]} {
    %val = "tf.StatefulPartitionedCall"(%arg0) {config = "", config_proto = "", executor_type = "", f = @f_callee} : (tensor<!tf_type.resource<tensor<f32>>>) -> (tensor<f32>)
    func.return %val : tensor<f32>
  }

        }
      }
    }
  }
  return LogicalResult::success();
}

struct Callers {
  TF::StatefulPartitionedCallOp forward;
  TF::StatefulPartitionedCallOp core_tpu;
  TF::StatefulPartitionedCallOp backward;
  TF::StatefulPartitionedCallOp non_tpu;
};

template <typename InputContainer>
std::vector<Type> GetValueTypes(const InputContainer& input) {

        }
      }
    }
    ops_to_process.swap(new_ops);
    new_ops.clear();
  }
}

TF::StatefulPartitionedCallOp MakeFuncCaller(
    mlir::OpBuilder& builder, const Location& loc, func::FuncOp func,
    const llvm::SetVector<Value>& operands) {
  // Constructs a tf.StatefulPartitionedCall to the function provided in 'func'
  // using the operands in 'operands'. Assumes the insertion point on builder is

//         dtype: DT_FLOAT
//         tensor_shape: {
//         }
//       }
//     }
//     outputs: {
//       key  : "output_0"
//       value: {
//         name : "StatefulPartitionedCall:0"
//         dtype: DT_FLOAT
//         tensor_shape: {
//         }
//       }
//     }
//     method_name: "tensorflow/serving/predict"
//   }
// }

  auto result = module.walk([&](Operation* op) {
    return llvm::isa_and_nonnull<mlir::TF::StatefulPartitionedCallOp>(op)
               ? mlir::WalkResult::interrupt()
               : mlir::WalkResult::advance();
  });
  if (result.wasInterrupted()) {
    // StatefulPartitionedCall ops are not supported by the tflite runtime.
    mlir::TFL::AttachErrorCode(
        module.emitError(

  // CHECK: [[results_0:%.*]]:2 = tfrt.call @callee([[chain]]
  // CHECK-SAME: (!tfrt.chain, !tfrt_fallback.tf_tensor) -> (!tfrt.chain, !tfrt_fallback.tf_tensor)
  %1 = "tf.StatefulPartitionedCall"(%0) {config = "", config_proto = "", executor_type = "", f = @callee} : (tensor<i32>) -> (tensor<i32>)
  // CHECK-NEXT: [[results_1:%.*]]:2 = tfrt.call @callee([[chain]]

  func.func @serving_default(%arg0: tensor<8xf32> {tf_saved_model.index_path = ["input"]}) -> (tensor<8xf32> {tf_saved_model.index_path = ["output"]})
    attributes {tf.entry_function = {control_outputs = "", inputs = "", outputs = "StatefulPartitionedCall:0"}, tf_saved_model.exported_names = ["serving_default"]} {
    %cst_0 = "tf.Const"() {device = "", value = dense<1.0> : tensor<8xf32>} : () -> tensor<8xf32>

      result.replaceAllUsesWith(if_op.getOperand(then_aliased_input + 1));
    }
  }
  if_op.replaceAllUsesWith(new_if);
  if_op.erase();
  return success();
}

// Handles stack usage by a tf.StatefulPartitionedCall or a tf.PartitionedCall.
// It will first check if the callee was previously handled, and try to reuse
// that result if so. Otherwise, it will clone and convert the callee function,

  // CHECK: [[results_0:%.*]] = call @callee(
  // CHECK-SAME: (!tf_mlrt.tensor) -> !tf_mlrt.tensor
  %1 = "tf.StatefulPartitionedCall"(%0) {config = "", config_proto = "", executor_type = "", f = @callee} : (tensor<i32>) -> (tensor<i32>)
  // CHECK-NEXT: [[results_1:%.*]] = call @callee(
  // CHECK-SAME: (!tf_mlrt.tensor) -> !tf_mlrt.tensor

                                                arg_buffer_size_is_fixed(i)};
    }
  }
  UpdateFuncType(func);
}

// Holds information about a decomposed callee function for
// PartitionedCall/StatefulPartitionedCall.
struct PartitionedCallDecompositionInfo {
  bool signature_change;
  func::FuncOp decomposed_callee;
  llvm::SmallDenseMap<int64_t, int64_t> buffer_arg_to_size_arg;