reformat

docs/gen_ref_pages.py

@@ -16,12 +16,15 @@
 
     parts = tuple(module_path.parts)
 
+    # Only document classes imported in __init__.py
     if parts[-1] == "__init__":
         parts = parts[:-1]
         doc_path = doc_path.with_name("index.md")
         full_doc_path = full_doc_path.with_name("index.md")
-    elif parts[-1].startswith("_"):
+    else:
         continue
+    # elif parts[-1].startswith("_"):
+    #    continue
 
     nav[parts] = doc_path.as_posix()
 
@@ -32,4 +35,4 @@
     mkdocs_gen_files.set_edit_path(full_doc_path, path.relative_to(root))
 
 with mkdocs_gen_files.open("reference/SUMMARY.md", "w") as nav_file:
-    nav_file.writelines(nav.build_literate_nav())
+    nav_file.writelines(nav.build_literate_nav(1))

src/rs_distributions/distributions/folded_normal.py

@@ -133,7 +133,8 @@ def pdf(self, value):
 
     def rsample(self, sample_shape=torch.Size()):
         """
-        Generate differentiable random samples from the Folded Normal distribution
+        Generate differentiable random samples from the Folded Normal distribution.
+        Gradients are implemented using implicit reparameterization (https://arxiv.org/abs/1805.08498).
 
         Args:
             sample_shape (torch.Size, optional): The shape of the samples to generate.