From 5526fde28b81445bc416504d93c3a1ca931e3d0d Mon Sep 17 00:00:00 2001 From: Maxim Ziatdinov Date: Sun, 20 Oct 2024 19:14:02 -0700 Subject: [PATCH] Update docstrings --- neurobayes/flax_nets/deterministic_nn.py | 3 +-- neurobayes/models/partial_bnn.py | 2 +- neurobayes/models/partial_bnn_heteroskedastic.py | 2 +- neurobayes/models/partial_dkl.py | 2 +- 4 files changed, 4 insertions(+), 5 deletions(-) diff --git a/neurobayes/flax_nets/deterministic_nn.py b/neurobayes/flax_nets/deterministic_nn.py index c8e5cc0..d1da3e6 100644 --- 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 index c2ad564..fb60c74 100644 --- 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 index 40d6f10..758d6f9 100644 --- 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 index cf12b3d..fe24dcf 100644 --- 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,