Merge pull request #88 from scikit-learn-contrib/joss-feedback

Joss feedback @mbsuraj

Author: EmanuelSommer

CONTRIBUTING.md

@@ -0,0 +1,43 @@
+# Contributing to bde
+
+Thanks for your interest in improving `bde`. We welcome bug reports,
+documentation fixes, examples, tests, and focused code contributions.
+
+Please use GitHub as the default place for project communication:
+
+- Issues: https://github.com/scikit-learn-contrib/bde/issues
+- Pull requests: https://github.com/scikit-learn-contrib/bde/pulls
+
+## Contributing changes
+
+For small fixes, open a pull request directly. For larger changes or API changes,
+please open an issue first so the approach can be discussed.
+
+Recommended workflow:
+
+1. Fork the repository.
+2. Create a branch for your change.
+3. Install the development environment as outlined in the README and do not forget to also install the `pre-commit` hooks.
+4. Add tests or documentation when relevant.
+5. Run the relevant checks (lint,tests).
+6. Open a pull request on GitHub and link any related issue.
+
+## Reporting issues
+
+Please report bugs, feature requests, and documentation problems through GitHub
+Issues. Before opening a new issue, check whether a similar issue already exists.
+
+For bug reports, include:
+
+- A minimal reproducible example.
+- The full traceback or error message.
+- Your operating system, Python version, and relevant package versions.
+- The expected behavior and the behavior you observed.
+
+## Seeking support
+
+Please use GitHub Issues for usage questions and support requests. This keeps
+answers visible to other users and helps the project maintain a searchable record.
+
+Direct email to the maintainers should be reserved for rare cases that are not
+appropriate for a public GitHub issue, such as sensitive or private matters.

README.md

@@ -63,6 +63,8 @@ pixi run python -m examples.example
 Pixi ensures the correct JAX, CUDA (when needed), and scikit-learn versions are
 selected automatically. See `pixi.lock` for channel and platform details.
 
+For contributor tooling, use the `dev` Pixi environment. It combines the
+documentation, linting, and test environments.
 
 
 Example Usage

bde/bde.py

@@ -843,7 +843,7 @@ def __init__(
         seed : int, default=0
             Shared PRNG seed for member initialisation and sampling.
         loss : BaseLoss | None
-            Custom training loss; defaults to :class:`bde.loss.GaussianNLL`.
+            Custom training loss; defaults to :class:`bde.loss.CategoricalCrossEntropy`.
         activation : str, default='relu'
             Activation function applied to each hidden layer.
         epochs : int, default=20

bde/loss/loss.py

@@ -26,8 +26,8 @@ def name(self) -> str:
 
     def __call__(self, preds: ArrayLike, y_true: ArrayLike):
         mu = preds[..., 0:1]
-        rmse = (mu - y_true) ** 2
-        return jnp.mean(rmse)
+        squared_error = (mu - y_true) ** 2
+        return jnp.sqrt(jnp.mean(squared_error))
 
 
 class GaussianNLL(BaseLoss):

paper/paper.bib

@@ -5,6 +5,14 @@ @inproceedings{sommer2025mile
     year={2025},
 }
 
+@inproceedings{sommer2026smile,
+  title  = {{Can Microcanonical Langevin Dynamics Leverage Mini-Batch Gradient Noise?}},
+  author = {Emanuel Sommer and Kangning Diao and Jakob Robnik and Uroš Seljak and David R{\"u}gamer},
+  booktitle={Forty-third International Conference on Machine Learning},
+  year={2026},
+  publisher={PMLR},
+}
+
 @inproceedings{sommer2024connecting,
   title={Connecting the Dots: Is Mode-Connectedness the Key to Feasible Sample-Based Inference in Bayesian Neural Networks?},
   author={Sommer, Emanuel and Wimmer, Lisa and Papamarkou, Theodore and Bothmann, Ludwig and Bischl, Bernd and R{\"u}gamer, David},
@@ -21,6 +29,38 @@ @InProceedings{papamarkou2024position
   publisher = {PMLR}
 }
 
+@article{pymc2023,
+  title = {{PyMC}: A Modern and Comprehensive Probabilistic Programming Framework in {P}ython},
+  author = {Oriol Abril-Pla and Virgile Andreani and Colin Carroll and Larry Dong and Christopher J. Fonnesbeck and Maxim Kochurov and Ravin Kumar and Junpeng Lao and Christian C. Luhmann and Osvaldo A. Martin and Michael Osthege and Ricardo Vieira and Thomas Wiecki and Robert Zinkov },
+  journal = {{PeerJ} Computer Science},
+  volume = {9},
+  number = {e1516},
+  doi = {10.7717/peerj-cs.1516},
+  year = {2023}
+}
+
+@article{phan2019numpyro,
+  title={Composable Effects for Flexible and Accelerated Probabilistic Programming in NumPyro},
+  author={Phan, Du and Pradhan, Neeraj and Jankowiak, Martin},
+  journal={arXiv preprint arXiv:1912.11554},
+  year={2019}
+}
+@inproceedings{daxberger2021laplace,
+  title = {Laplace {{Redux}} {\textendash} {{Effortless Bayesian Deep Learning}}},
+  booktitle = {35th {{Conference}} on {{Neural Information Processing Systems}} ({{NeurIPS}} 2021)},
+  author = {Daxberger, Erik and Kristiadi, Agustinus and Immer, Alexander and Eschenhagen, Runa and Bauer, Matthias and Hennig, Philipp},
+  year = {2021},
+  keywords = {approximate\_inference,bnn,gaussian\_posterior,laplace,last\_layer,uncertainty\_quantification}
+}
+
+@inproceedings{
+    duffield2025scalable,
+    title={Scalable Bayesian Learning with posteriors},
+    author={Samuel Duffield and Kaelan Donatella and Johnathan Chiu and Phoebe Klett and Daniel Simpson},
+    booktitle={The Thirteenth International Conference on Learning Representations},
+    year={2025}
+}
+
 @misc{cabezas2024blackjax,
       title={{BlackJAX: Composable Bayesian inference in JAX}},
       author={Alberto Cabezas and Adrien Corenflos and Junpeng Lao and Rémi Louf},
@@ -92,4 +132,3 @@ @misc{misc_bike_sharing_dataset_275
   year         = {2013},
   howpublished = {UCI Machine Learning Repository}
 }
-

paper/paper.md

@@ -47,10 +47,11 @@ The workflow of `bde` implements the two-stage BDE inference process of MILE. Fi
 # Software design
 
 Because optimization and sampling across ensemble members are independent, `bde` exploits JAX’s parallelization and just-in-time compilation to scale efficiently across CPUs, GPUs, and TPUs. Given new test data, the package approximates the posterior predictive, enabling point predictions, credible intervals, coverage estimates, and other uncertainty metrics through a unified interface.
+Further, this work currently targets the well-studied and empirically validated full-batch, fully connected-network setting for tabular data also studied in MILE [@sommer2025mile]. This makes it particularly suitable for the chosen scikit-learn framework. Stochastic-gradient variants such as SMILE [@sommer2026smile] require different algorithmic and software design choices, yet will be considered for future extensions.
 
 # State of the field
 
-Reliable uncertainty quantification (UQ) is increasingly viewed as a critical component of modern machine learning systems, and BDL provides a principled framework for achieving it [@papamarkou2024position]. While several libraries support optimization-based approaches such as variational inference or classical Bayesian modeling, accessible tools for sampling-based inference in Bayesian neural networks remain scarce. Existing probabilistic programming frameworks offer MCMC but require substantial manual configuration to achieve competitive performance on neural network models.
+Reliable uncertainty quantification (UQ) is increasingly viewed as a critical component of modern machine learning systems, and BDL provides a principled framework for achieving it [@papamarkou2024position]. While several libraries support optimization-based approaches such as variational inference [@duffield2025scalable], Laplace approximations [@daxberger2021laplace], or classical Bayesian modeling [@pymc2023], accessible tools for sampling-based inference in Bayesian neural networks remain scarce. Existing probabilistic programming and inference frameworks such as NumPyro [@phan2019numpyro] and `blackjax` [@cabezas2024blackjax] offer MCMC building blocks but require substantial manual configuration to achieve competitive Bayesian neural-network workflows.
 
 # Statement of need
 
@@ -97,7 +98,7 @@ Classification follows analogously using `BdeClassifier`.
 
 # Regression benchmark
 
-We provide a small benchmark of `bde` on the `airfoil` [@Dua_2019] and the `bikesharing` [@misc_bike_sharing_dataset_275] datasets. We report mean predictive performance (RMSE), UQ metrics (NLL in the distributional and mean regression formulation), reported as mean ± standard deviation over 5 independent runs. The results show competitive out-of-the-box performance of BDE especially in UQ with its native distributional regression capability.
+We provide a small benchmark of `bde` on the `airfoil` [@Dua_2019] and the `bikesharing` [@misc_bike_sharing_dataset_275] datasets. We report mean predictive performance (RMSE), UQ metrics (NLL in the distributional and mean regression formulation), reported as mean ± standard deviation over 5 independent model seeds and fixed train/test splits. The results show competitive out-of-the-box performance of BDE especially in UQ with its native distributional regression capability.
 
 | `airfoil` | RMSE                | NLL (distr. regr.)   | NLL (mean regr.)     |
 | ------------- | ------------------- | -------------------- | -------------------- |