Improve the type hints of _get_params_or_grads() (#170)

Summary:
Pull Request resolved: #170

Figure out how to type hint when there is a default value of `get_grad=False` in `_get_params_or_grads()`.

Reviewed By: gajjanag

Differential Revision: D74560386

fbshipit-source-id: 691f4dd5a2fa4ed8875d1942801cdaa24afe669a

Author: tsunghsienlee

distributed_shampoo/utils/shampoo_distributor.py

@@ -109,10 +109,6 @@ def _construct_composable_block_ids(
         """
         return (param_index, f"block_{block_index}")
 
-    @overload
-    @torch.no_grad()
-    def _get_params_or_grads(self) -> Iterable[Tensor]: ...
-
     @overload
     @torch.no_grad()
     def _get_params_or_grads(
@@ -121,7 +117,9 @@ def _get_params_or_grads(
 
     @overload
     @torch.no_grad()
-    def _get_params_or_grads(self, get_grad: Literal[False]) -> Iterable[Tensor]: ...
+    def _get_params_or_grads(
+        self, get_grad: Literal[False] = False
+    ) -> Iterable[Tensor]: ...
 
     @torch.no_grad()
     def _get_params_or_grads(self, get_grad: bool = False) -> Iterable[Tensor | None]:

distributed_shampoo/utils/shampoo_fully_shard_distributor.py

@@ -29,10 +29,6 @@ class FullyShardDistributor(Distributor):
 
     """
 
-    @overload
-    @torch.no_grad()
-    def _get_params_or_grads(self) -> Iterable[Tensor]: ...
-
     @overload
     @torch.no_grad()
     def _get_params_or_grads(
@@ -41,7 +37,9 @@ def _get_params_or_grads(
 
     @overload
     @torch.no_grad()
-    def _get_params_or_grads(self, get_grad: Literal[False]) -> Iterable[Tensor]: ...
+    def _get_params_or_grads(
+        self, get_grad: Literal[False] = False
+    ) -> Iterable[Tensor]: ...
 
     @torch.no_grad()
     def _get_params_or_grads(self, get_grad: bool = False) -> Iterable[Tensor | None]:

distributed_shampoo/utils/shampoo_hybrid_shard_distributor.py

@@ -203,10 +203,6 @@ def __init__(
             comms_group_rank=comms_group_rank,
         )
 
-    @overload
-    @torch.no_grad()
-    def _get_params_or_grads(self) -> Iterable[Tensor]: ...
-
     @overload
     @torch.no_grad()
     def _get_params_or_grads(
@@ -215,7 +211,9 @@ def _get_params_or_grads(
 
     @overload
     @torch.no_grad()
-    def _get_params_or_grads(self, get_grad: Literal[False]) -> Iterable[Tensor]: ...
+    def _get_params_or_grads(
+        self, get_grad: Literal[False] = False
+    ) -> Iterable[Tensor]: ...
 
     @torch.no_grad()
     def _get_params_or_grads(self, get_grad: bool = False) -> Iterable[Tensor | None]: