Better type hint of instantiate_preconditioner_config()

Summary: `instantiate_preconditioner_config()` only returns `None` when encoutnering `PreconditionerComputationType.None` as input, and this type hint enhancements make the intention more clear.

Differential Revision: D81278578

Author: tsunghsienlee

distributed_shampoo/examples/trainer_utils.py

@@ -17,7 +17,7 @@
 from functools import partial
 from operator import attrgetter
 from pathlib import Path
-from typing import overload, Type
+from typing import Literal, overload, Type
 
 import numpy as np
 
@@ -406,11 +406,6 @@ def instantiate_optimizer(
             weight_decay=weight_decay,
         )
     elif optimizer_type == OptimizerType.DISTRIBUTED_SHAMPOO:
-        assert (
-            preconditioner_config := instantiate_preconditioner_config(
-                preconditioner_computation_type=preconditioner_computation_type,
-            )
-        ) is not None
         optimizer_cls = partial(
             DistributedShampoo,
             betas=betas,
@@ -431,14 +426,32 @@ def instantiate_optimizer(
                 grafting_epsilon=grafting_epsilon,
             ),
             distributed_config=distributed_config,
-            preconditioner_config=preconditioner_config,
+            preconditioner_config=instantiate_preconditioner_config(
+                preconditioner_computation_type=preconditioner_computation_type,
+            ),
         )
     else:
         raise ValueError(f"Invalid OptimizerType {optimizer_type}!")
 
     return optimizer_cls(parameters, lr=lr)
 
 
+@overload
+def instantiate_preconditioner_config(
+    preconditioner_computation_type: Literal[PreconditionerComputationType.NONE],
+    grafting_beta2: float = ...,
+    grafting_epsilon: float = ...,
+) -> None: ...
+
+
+@overload
+def instantiate_preconditioner_config(
+    preconditioner_computation_type: PreconditionerComputationType,
+    grafting_beta2: float = ...,
+    grafting_epsilon: float = ...,
+) -> PreconditionerConfig: ...
+
+
 def instantiate_preconditioner_config(
     preconditioner_computation_type: PreconditionerComputationType,
     grafting_beta2: float = 1.0,