Variational Neural Networks implementation in Pytorch and JAX[Formula presented]

Illia Oleksiienko*, Dat Thanh Tran, Alexandros Iosifidis

*Corresponding author for this work

    Research output: Contribution to journal/Conference contribution in journal/Contribution to newspaperJournal articleResearchpeer-review


    Bayesian Neural Networks consider a distribution over the network's weights, which provides a tool to estimate the uncertainty of a neural network by sampling different models for each input. Variational Neural Networks (VNNs) consider a probability distribution over each layer's outputs and generate parameters for it with the corresponding sub-layers. We provide two Python implementations of VNNs with PyTorch and JAX machine learning libraries that ensure reproducibility of the experimental results and allow implementing uncertainty estimation methods easily in other projects.

    Original languageEnglish
    Article number100431
    JournalSoftware Impacts
    Publication statusPublished - Nov 2022


    • Bayesian deep learning
    • Bayesian Neural Networks
    • JAX
    • PyTorch
    • Uncertainty estimation


    Dive into the research topics of 'Variational Neural Networks implementation in Pytorch and JAX[Formula presented]'. Together they form a unique fingerprint.

    Cite this