add docs (#7848)

Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

python/oneflow/nn/modules/fused_mlp.py

@@ -77,6 +77,9 @@ def __init__(
         self.reset_parameters()
 
     def add_parameters(self) -> None:
+        """Register parameter in FusedMLP module. 
+
+        """
         if self.hidden_layer_num != 0:
             # First layer.
             self.register_parameter(
@@ -129,18 +132,33 @@ def add_parameters(self) -> None:
             )
 
     def weight(self, i):
+        """Returns the ith weight. 
+
+        """
         return getattr(self, f"weight_{i}")
 
     def weights(self):
+        """Returns the weight list in FusedMLP module. 
+
+        """
         return [self.weight(i) for i in range(self.hidden_layer_num + 1)]
 
     def bias(self, i):
+        """Return the ith bias. 
+
+        """
         return getattr(self, f"bias_{i}")
 
     def biases(self):
+        """Returns the bias list in FusedMLP module. 
+
+        """
         return [self.bias(i) for i in range(self.hidden_layer_num + 1)]
 
     def reset_parameters(self) -> None:
+        """Reset the parameters in FusedMLP module. 
+
+        """
         for layer_idx in range(self.hidden_layer_num + 1):
             flow.nn.init.kaiming_uniform_(self.weight(layer_idx), a=math.sqrt(5))
             (fan_in, _) = _calculate_fan_in_and_fan_out(self.weight(layer_idx))

python/oneflow/nn/optimizer/adamw.py

@@ -269,4 +269,7 @@ def _generate_conf_for_graph(self, train_conf, vars_conf):
 
     @property
     def support_sparse(self):
+        """Whether AdamW Optimizer support sparse update. 
+
+        """
         return True

python/oneflow/nn/optimizer/lambda_lr.py

@@ -94,6 +94,9 @@ def load_state_dict(self, state_dict):
                 self.lr_lambdas[idx].__dict__.update(fn)
 
     def step(self):
+        """Performs a single learning rate schedule step.
+
+        """
         self.last_step += 1
         lrs = []
         for (lmbda, base_lr) in zip(self.lr_lambdas, self.base_lrs):

python/oneflow/nn/optimizer/optimizer.py

@@ -232,6 +232,14 @@ def pack_group(group):
         }
 
     def step(self, closure: Union[Callable, None] = None) -> Union[Tensor, None]:
+        """Performs a single optimization step (parameter update).
+
+        Args:
+            closure (Union[Callable, None], optional): A closure that reevaluates the model and returns the loss. Optional for most optimizers.
+
+        Returns:
+            Union[Tensor, None]: The loss. 
+        """
         raise NotImplementedError()
 
     def clip_grad(self):
@@ -329,6 +337,9 @@ def _generate_grad_clip_conf_for_optim_conf(self, param_group, optimizer_conf):
 
     @property
     def support_sparse(self):
+        """Whether the Optimizer support sparse update. 
+
+        """
         return False
 
     def _check_variables_in_graph(self, vars_conf):

python/oneflow/nn/optimizer/reduce_lr_on_plateau.py

@@ -121,7 +121,7 @@ def __init__(
         self._reset()
 
     def step(self, metrics):
-        """Step forward once.
+        """Performs a single learning rate schedule step.
 
         Arguments:
             metrics (float): a metrics quantity of Measuring the effect of model training.
@@ -149,9 +149,15 @@ def step(self, metrics):
 
     @property
     def in_cooldown(self):
+        """Whether the learning rate scheduler in cooldown phase. 
+
+        """
         return self.cooldown_counter > 0
 
     def is_better(self, a, best):
+        """Whether the metric has improvement. 
+        
+        """
         if self.mode == "min" and self.threshold_mode == "rel":
             rel_epsilon = 1.0 - self.threshold
             return a < best * rel_epsilon

python/oneflow/nn/optimizer/sgd.py

@@ -128,6 +128,12 @@ def __init__(
         )
 
     def step(self, closure: Callable = None):
+        """Performs a single optimization step.
+
+        Args:
+            closure (callable, optional): A closure that reevaluates the model
+                and returns the loss.
+        """
         with flow.no_grad():
             loss = None
             if closure is not None:
@@ -187,4 +193,7 @@ def _generate_conf_for_graph(self, train_conf, vars_conf):
 
     @property
     def support_sparse(self):
+        """Whether SGD Optimizer support sparse update. 
+
+        """
         return True