Update docstrings

ziatdinovmax · Oct 21, 2024 · 5526fde · 5526fde
1 parent 97e4170
commit 5526fde
Show file tree

Hide file tree

Showing 4 changed files with 4 additions and 5 deletions.
diff --git a/neurobayes/flax_nets/deterministic_nn.py b/neurobayes/flax_nets/deterministic_nn.py
@@ -16,8 +16,7 @@ class TrainState(train_state.TrainState):
 class DeterministicNN:
 
     """
-    Class for training and predicting with deterministic (as in non-stochastic)
-    neural network.
+    Class for training and predicting with deterministic neural network.
 
     Args:
         architecture: a Flax model

diff --git a/neurobayes/models/partial_bnn.py b/neurobayes/models/partial_bnn.py
@@ -29,7 +29,7 @@ class PartialBNN(BNN):
     SPLITTERS = {
         FlaxMLP: split_mlp,
         FlaxConvNet: split_convnet,
-        # Add more network types and their splitters here
+        # More network types and their splitters TBA
     }
 
     def __init__(self,

diff --git a/neurobayes/models/partial_bnn_heteroskedastic.py b/neurobayes/models/partial_bnn_heteroskedastic.py
@@ -18,7 +18,7 @@ class HeteroskedasticPartialBNN(HeteroskedasticBNN):
     SPLITTERS = {
         FlaxMLP2Head: split_mlp2head,
         FlaxConvNet2Head: split_convnet2head,
-        # Add more network types and their splitters here
+        # More network types and their splitters TBA
     }
     def __init__(self,
                  deterministic_nn: Type[flax.linen.Module],

diff --git a/neurobayes/models/partial_dkl.py b/neurobayes/models/partial_dkl.py
@@ -23,7 +23,7 @@ class PartialDKL(DKL):
     SPLITTERS = {
         FlaxMLP: split_mlp,
         FlaxConvNet: split_convnet,
-        # Add more network types and their splitters here
+        # More network types and their splitters TBA
     }
 
     def __init__(self,