chore: update default arguments

metaopt · XuehaiPan · Oct 8, 2022 · Sep 29, 2022 · Sep 30, 2022 · Oct 7, 2022
commit 891d6d09b6f491d13ce0a42157f002b29b63fc66
diff --git a/torchopt/distributed/api.py b/torchopt/distributed/api.py
@@ -365,7 +365,7 @@ def remote_sync_call(
 
 
 def parallelize_async(
-    partitioner: Partitioner = batch_partitioner,
+    partitioner: Optional[Partitioner] = None,
     reducer: Optional[Callable[[Iterable[T]], U]] = None,
     timeout: Optional[float] = UNSET_RPC_TIMEOUT,
 ) -> Callable[[Callable[..., T]], Callable[..., Union[Future[List[T]], Future[U]]]]:
@@ -379,14 +379,20 @@ def parallelize_async(
         partitioner (Partitioner, optional): A partitioner that partitions the arguments to multiple
             workers. Defaults to :func:`batch_partitioner`.
         reducer (Callable[[Iterable[T]], U], optional): A reducer that reduces the results from
-            multiple workers. Defaults to :data:`None`.
+            multiple workers. Defaults to :func:`mean_reducer` if the ``partitioner`` is not
+            specified, i.e., :func:`batch_partitioner`. Otherwise, it defaults to :data:`None`.
         timeout (float, optional): The timeout for the RPC call. Defaults to
             :data:`rpc.api.UNSET_RPC_TIMEOUT`.
 
     Returns:
         The decorated function.
     """
 
+    if partitioner is None:
+        partitioner = batch_partitioner
+        if reducer is None:
+            reducer = mean_reducer
+
     def wrapper(func: Callable[..., T]) -> Callable[..., Union[Future[List[T]], Future[U]]]:
         @functools.wraps(func)
         def wrapped(*args: Any, **kwargs: Any) -> Union[Future[List[T]], Future[U]]:
@@ -418,7 +424,7 @@ def wrapped(*args: Any, **kwargs: Any) -> Union[Future[List[T]], Future[U]]:
 
 
 def parallelize(
-    partitioner: Partitioner = batch_partitioner,
+    partitioner: Optional[Partitioner] = None,
     reducer: Optional[Callable[[Iterable[T]], U]] = None,
     timeout: Optional[float] = UNSET_RPC_TIMEOUT,
 ) -> Callable[[Callable[..., T]], Callable[..., Union[List[T], U]]]:
@@ -430,14 +436,20 @@ def parallelize(
         partitioner (Partitioner, optional): A partitioner that partitions the arguments to multiple
             workers. Defaults to :func:`batch_partitioner`.
         reducer (Callable[[Iterable[T]], U], optional): A reducer that reduces the results from
-            multiple workers. Defaults to :data:`None`.
+            multiple workers. Defaults to :func:`mean_reducer` if the ``partitioner`` is not
+            specified, i.e., :func:`batch_partitioner`. Otherwise, it defaults to :data:`None`.
         timeout (float, optional): The timeout for the RPC call. Defaults to
             :data:`rpc.api.UNSET_RPC_TIMEOUT`.
 
     Returns:
         The decorated function.
     """
 
+    if partitioner is None:
+        partitioner = batch_partitioner
+        if reducer is None:
+            reducer = mean_reducer
+
     def wrapper(func: Callable[..., T]) -> Callable[..., Union[List[T], U]]:
         @functools.wraps(func)
         def wrapped(*args: Any, **kwargs: Any) -> Union[List[T], U]: