What's Changed
Add an option to specify which layers in the provided architecture will be treated as probabilistic. For example,
# Initialize NN architecture
architecture = nb.FlaxMLP(hidden_dims = [64, 32, 16, 8], target_dim=1)
# Make the first and output layers probabilistic and the rest deterministic
probabilistic_layer_names = ['Dense0', 'Dense4']
# Intitalize and train a PBNN model
model = nb.PartialBNN(architecture, probabilistic_layer_names=probabilistic_layer_names)
model.fit(X_measured, y_measured, num_warmup=1000, num_samples=1000)Full Changelog: 0.0.7...0.0.9