elif node.id in TFR_BUILTINS:
      val = node.id
      lookup_type = anno.getanno(node, anno.Static.TYPES, types.FunctionType)
    else:
      op_def, _ = self._op_defs.lookup(node.id)
      val = op_def.name
      lookup_type = anno.getanno(node, anno.Static.TYPES, types.FunctionType)
    type_ = self._get_inferred_type(node, lookup_type)
    return val, type_

  def visit_Return(self, node):

      llvm::StringMap<Attribute>* derived_attrs) const;

  // Uses the collected attribute values to derive all the output types.
  LogicalResult DeriveOutputTypes(Location loc, FunctionType signature,
                                  const llvm::StringMap<Attribute>& attrs,
                                  SmallVectorImpl<Type>* output_types) const;

  llvm::SmallVector<Type> input_types =
      TypesFromValues(subgraph.FuncArguments());
  llvm::SmallVector<Type> return_types =
      TypesFromValues(subgraph.FuncOutputs());

  FunctionType function_type =
      builder.getFunctionType(input_types, return_types);

  std::string function_name = absl::StrCat("func_", subgraph.subgraph_id_);

  builder.setInsertionPointToEnd(&module_op.getBodyRegion().front());

  auto filename_input_type = RankedTensorType::get(
      /*shape=*/{}, /*elementType=*/builder.getType<TF::StringType>());

  FunctionType func_type = builder.getFunctionType(
      /*inputs=*/{filename_input_type}, /*results=*/{});
  auto save_func = builder.create<func::FuncOp>(
      NameLoc::get(builder.getStringAttr(kTfQuantSaveFuncName)),

                                    PatternRewriter& rewriter,
                                    Block* insertion_point) {
  rewriter.setInsertionPointToStart(insertion_point);

  auto ftype = FunctionType::get(rewriter.getContext(), {type}, {type});
  auto name = rewriter.getStringAttr("gelu_decomp");
  func::FuncOp new_func = rewriter.create<func::FuncOp>(
      insertion_point->front().getLoc(), name, ftype);

    new_float_func_args[weight_operand_idx] = reshape_op;
    new_float_func.getArgument(weight_operand_idx).setType(new_shape);
    new_float_func.setType(FunctionType::get(
        getContext(), TypeRange{ValueRange{new_float_func_args}},
        new_float_func.getResultTypes()));
    symbol_table.insert(new_float_func);

		return types.NewTuple(types.NewParam(token.NoPos, pkg, "", taa))

	case '(':
		params, _ := p.parseParamList(pkg)
		return params

	default:
		return nil
	}
}

// FunctionType = ParamList ResultList .
func (p *parser) parseFunctionType(pkg *types.Package, nlist []any) *types.Signature {
	t := new(types.Signature)
	p.update(t, nlist)

	params, isVariadic := p.parseParamList(pkg)

  xla_input_shapes->clear();

  mlir::func::FuncOp main_func =
      module.lookupSymbol<mlir::func::FuncOp>("main");
  TF_RET_CHECK(main_func != nullptr) << "No main function found";
  mlir::FunctionType func_type = main_func.getFunctionType();

  int num_args = func_type.getNumInputs();
  xla_input_shapes->reserve(num_args);

  std::vector<xla::Shape> individual_arg_shapes;

  // Note, for a while loop body function, the operand types and result types
  // are identical.
  auto body_func_type = mlir::FunctionType::get(
      &getContext(), new_while_operand_types, new_while_operand_types);
  auto cond_func_type = mlir::FunctionType::get(
      &getContext(), new_while_operand_types, orig_cond_func.getResultTypes());
  func::FuncOp cond =

          .replaceAllUsesWith(
              builder.clone(*operand.getDefiningOp())->getResult(0));
    }
  }

  // Update/use input shapes for function.
  FunctionType func_type = func.getFunctionType();
  func.setType(FunctionType::get(func.getContext(), op.getOperandTypes(),
                                 func_type.getResults()));

  // Run shape inference on the function.
  if (failed(