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;

                                    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);

  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)),

    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);

  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(

  auto output_type = tensorflow::GetTypeFromTFTensorShape(
      output_shape,
      mlir::cast<RankedTensorType>(input_.getType()).getElementType());
  fused_func_op_.setType(mlir::FunctionType::get(fused_func_op_.getContext(),
                                                 input_types, output_type));
}

LogicalResult ConvertLSTMCellSimpleToFusedLSTM::RewriteFunc() {
  LogicalResult result = Initialize();

  // unranked tensors and could be ranked tensors in cases where result type
  // depends on an op with static output shape like tf.Const.
  absl::StatusOr<mlir::FunctionType> InferLibFunctionType(
      const FunctionBody& fbody);

  // Extracts arg and ret nodes from FunctionBody.
  void GetArgsAndRetsFromFunctionBody(
      const FunctionBody& fbody,

    Type branch_args_type[] = {input_handle.getType(), input_shape.getType(),
                               size_diff.getType(), size.getType()};
    Type branch_result_type[] = {result_type};
    auto func_type = FunctionType::get(rewriter.getContext(), branch_args_type,
                                       branch_result_type);

    // Create functions in a higher scope before restoring the insertion point.