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

  // Find any StatefulPartitionedCalls and inline their contents in this func.
  LogicalResult InlineCallsInFunc(func::FuncOp func,
                                  bool inline_all_funcs = false) {
    llvm::SetVector<Operation*> ops_to_erase;
    for (auto caller :