mlir::LogicalResult PromoteVarHandlesToArguments(
    func::FuncOp function, bool add_validation,
    llvm::SmallVectorImpl<std::string>* var_handle_shared_names) {
  Block& block = function.front();
  auto func_type = function.getFunctionType();

  auto func_arg_types = llvm::to_vector<4>(func_type.getInputs());
  llvm::SmallDenseMap<llvm::StringRef, int> var_arg_index_by_name;
  for (auto var_handle_op :

void RemovePlatformIndexArg(MLIRContext *ctx, func::FuncOp func_op) {
  // If there are multiple platforms, the first argument is reserved for
  // passing the platform index.
  FunctionType function_type = func_op.getFunctionType();
  ArrayRef<Type> new_input_types =
      function_type.getInputs().take_back(func_op.getNumArguments() - 1);
  func_op.setFunctionType(
      FunctionType::get(ctx, new_input_types, function_type.getResults()));

    llvm::function_ref<std::optional<Type>(int64_t)> arg_to_stack_type,
    llvm::function_ref<void(ArrayRef<BlockArgument>)> handle_new_size_vars =
        nullptr) {
  auto new_input_types = llvm::to_vector<8>(func.getFunctionType().getInputs());
  auto size_var_type = GetSizeVarType(OpBuilder(func));
  int64_t original_arg_count = new_input_types.size();
  for (int64_t i = 0; i < original_arg_count; ++i) {
    auto stack_type = arg_to_stack_type(i);

                                             custom_option_buffer))) {
    return failure();
  }

  auto tfl_fusable_op = builder.create<TFL::CustomOp>(
      func->getLoc(), func.getFunctionType().getResults(), func.getArguments(),
      custom_op_name, CustomOption(&builder, custom_option_buffer));
  builder.create<func::ReturnOp>(func->getLoc(), tfl_fusable_op.getResults());
  return success();
}

      module_body.push_front(other);
      other.addEntryBlock();
      OpBuilder builder(other.getRegion());
      auto call_op = builder.create<mlir::func::CallOp>(
          f.getLoc(), f.getFunctionType().getResults(), f.getSymName(),
          other.getRegion().getArguments());
      builder.create<mlir::func::ReturnOp>(f.getLoc(), call_op.getResults());
    }

    Unexport(f);
  }
}

    llvm::function_ref<std::optional<Type>(int64_t)> arg_to_buffer_type,
    llvm::function_ref<bool(int64_t)> arg_buffer_size_is_fixed) {
  auto new_input_types = llvm::to_vector<8>(func.getFunctionType().getInputs());
  int64_t original_arg_count = new_input_types.size();
  Location loc = func.getLoc();
  for (int64_t i = 0; i < original_arg_count; ++i) {
    auto buffer_type = arg_to_buffer_type(i);

    SmallVectorImpl<Value>* inputs, NamedAttrList* arg_attrs,
    llvm::StringMap<Attribute>* derived_attrs) const {
  for (const auto& operand :
       llvm::enumerate(signature.getFunctionType().getInputs())) {
    // If the index is larger than the operand number of the call_op, the
    // default value of the operand needs to be used.
    if (operand.index() >= call_op.getNumOperands()) {

    // functions.
    auto func_names = mlir::tf_saved_model::GetExportedNames(func);
    if (func_names.empty()) return mlir::WalkResult::advance();

    auto func_type = func.getFunctionType();

    // Here we walk through each arguments and find out the input/output names,
    // and input devices, variables used by this function.
    llvm::SmallVector<llvm::StringRef, 4> input_names;
    llvm::SmallVector<

}

void LegalizeJaxRandomPass::runOnOperation() {
  auto func = getOperation();
  if (!IsJaxRandomUniform(func) && !IsJaxRandomNormal(func)) return;
  auto result_tuple_ty =
      mlir::dyn_cast_or_null<TupleType>(func.getFunctionType().getResult(0));
  if (!result_tuple_ty) return;
  if (result_tuple_ty.size() != 1) return;
  auto result_ty = mlir::dyn_cast<ShapedType>(result_tuple_ty.getType(0));

  func.eraseBody();
  func.addEntryBlock();

void UpdateFuncType(func::FuncOp func) {
  Operation* terminator = func.front().getTerminator();
  auto return_types = llvm::to_vector<4>(terminator->getOperandTypes());

  FunctionType func_type = func.getFunctionType();
  if (llvm::ArrayRef(return_types) == func_type.getResults()) return;

  auto updated_type =
      FunctionType::get(func.getContext(), func_type.getInputs(), return_types);
  func.setType(updated_type);
}