if (input_shapes_attr.size() != func_op.getNumArguments()) {
    return func_op->emitError(
               "Number of arguments and 'tf._input_shapes' "
               "attribute size do not match. ")
           << "Num args: " << func_op.getNumArguments()
           << ", tf._input_shapes size: " << input_shapes_attr.size();
  }

  return success();
}

// Validates ModuleOp. Returns `failure` if the module op is invalid.

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>

  // 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) {

    const std::vector<std::optional<std::vector<int>>>& input_shapes,
    const std::vector<std::string>& output_arrays,
    const std::vector<std::string>& control_output_arrays,
    const GraphdefToMlirOptions& import_options, mlir::MLIRContext* context) {
  auto module_or = GraphdefToMlirImport(
      input, input_arrays, input_dtypes, input_shapes, output_arrays,
      control_output_arrays, import_options, context);

        constant_pool.Get1DHostConstant(i), constant_pool.Get1DHostConstant(1));

    Output input_shape_i = ops::Slice(
        host_scope.WithOpName("input_shape_", i), input_shape,
        constant_pool.Get1DHostConstant(i), constant_pool.Get1DHostConstant(1));

    slice_size.push_back(ops::Sub(host_scope.WithOpName("slice_size_", i),

void HandleConv2DInput(TF::Conv2DOp conv2d, int64_t block_size) {
  auto input = conv2d.getInput();
  auto input_shape = mlir::cast<RankedTensorType>(input.getType()).getShape();
  SmallVector<int64_t, 4> transform_shape = {
      input_shape[0], input_shape[1] / block_size, input_shape[2] / block_size,
      input_shape[3] * block_size * block_size};
  auto transform_result_type =

    // Graphdef import path.
    module = tensorflow::LoadFromGraphdefOrMlirSource(
        input_file_name, input_mlir, use_splatted_constant, custom_opdefs,
        specs, debug_info_file, input_arrays, input_dtypes, input_shapes,
        output_arrays, control_output_arrays, &source_mgr, &context);
  }

  // If errors occur, the library call in the above already logged the error
  // message. So we can just return here.