Repository: google/flax

Issue:
Add reference documentation for FrozenDict on ReadTheDocs



Discussion:

Resolving patch:
diff --git a/flax/core/frozen_dict.py b/flax/core/frozen_dict.py
--- a/flax/core/frozen_dict.py
+++ b/flax/core/frozen_dict.py
@@ -120,9 +120,19 @@ def pop(self, key: K) -> Tuple['FrozenDict[K, V]', V]:
     return new_self, value
 
   def unfreeze(self) -> Dict[K, V]:
+    """Unfreeze this FrozenDict.
+
+    Returns:
+      An unfrozen version of this FrozenDict instance.
+    """
     return unfreeze(self)
 
-  def tree_flatten(self):
+  def tree_flatten(self) -> Tuple[Tuple[Dict[Any, Any]], Tuple[()]]:
+    """Flattens this FrozenDict.
+
+    Returns:
+      A flattened version of this FrozenDict instance.
+    """
     return (self._dict,), ()
 
   @classmethod
diff --git a/flax/training/checkpoints.py b/flax/training/checkpoints.py
--- a/flax/training/checkpoints.py
+++ b/flax/training/checkpoints.py
@@ -175,9 +175,12 @@ def restore_checkpoint(ckpt_dir,
 
   Sorts the checkpoint files naturally, returning the highest-valued
   file, e.g.:
-    ckpt_1, ckpt_2, ckpt_3 --> ckpt_3
-    ckpt_0.01, ckpt_0.1, ckpt_0.001 --> ckpt_0.1
-    ckpt_-1.0, ckpt_1.0, ckpt_1e5 --> ckpt_1e5
+
+  *  ``ckpt_1, ckpt_2, ckpt_3 --> ckpt_3``
+
+  *  ``ckpt_0.01, ckpt_0.1, ckpt_0.001 --> ckpt_0.1``
+
+  *  ``ckpt_-1.0, ckpt_1.0, ckpt_1e5 --> ckpt_1e5``
 
   Args:
     ckpt_dir: str: checkpoint file or directory of checkpoints to restore from.
@@ -252,7 +255,7 @@ def convert_pre_linen(params):
   submodule class. With Linen this behavior has changed to keep separate
   submodule counts per module class.
 
-  Consider the following module:
+  Consider the following module::
 
     class Model(nn.Module):
       @nn.compact
@@ -262,26 +265,28 @@ def __call__(self, x):
         return x
 
   In pre-Linen the resulting params would have had the structure:
-    {'Conv_0': { ... }, 'Dense_1': { ... } }
+
+    ``{'Conv_0': { ... }, 'Dense_1': { ... } }``
 
   With Linen the resulting params would instead have had the structure:
-    {'Conv_0': { ... }, 'Dense_0': { ... } }
 
-  To convert from pre-Linen format to Linen simply call:
+    ``{'Conv_0': { ... }, 'Dense_0': { ... } }``
+
+  To convert from pre-Linen format to Linen simply call::
 
     params = convert_pre_linen(pre_linen_params)
 
   Note that you can also use this utility to convert pre-Linen collections
   because they're following the same module naming. Note though that collections
   were "flat" in pre-Linen and first need to be unflattened before they can be
-  used with this function:
+  used with this function::
 
     batch_stats = convert_pre_linen(flax.traverse_util.unflatten_dict({
         tuple(k.split('/')[1:]): v
         for k, v in pre_linen_model_state.as_dict().items()
     }))
 
-  Then Linen variables can be defined from these converted collections:
+  Then Linen variables can be defined from these converted collections::
 
     variables = {'params': params, 'batch_stats': batch_stats}
 
diff --git a/flax/training/lr_schedule.py b/flax/training/lr_schedule.py
--- a/flax/training/lr_schedule.py
+++ b/flax/training/lr_schedule.py
@@ -58,12 +58,13 @@ def create_stepped_learning_rate_schedule(base_learning_rate, steps_per_epoch,
   by specified amounts at specified epochs. The steps are given as
   the `lr_sched_steps` parameter. A common ImageNet schedule decays the
   learning rate by a factor of 0.1 at epochs 30, 60 and 80. This would be
-  specified as:
-  [
-    [30, 0.1],
-    [60, 0.01],
-    [80, 0.001]
-  ]
+  specified as::
+
+    [
+      [30, 0.1],
+      [60, 0.01],
+      [80, 0.001]
+    ]
 
   This function also offers a learing rate warmup as per
   https://arxiv.org/abs/1706.02677, for the purpose of training with large


Test patch: