# TyXe: Pyro-based Bayesian neural nets for Pytorch

arXiv: https://arxiv.org/abs/2110.00276 · HTML: https://mlsanity.com/paper/2110.00276/

## Key facts

- Cited by 6 papers in the map (data through 14 Oct 2026).
- Builds on 26 of its 53 references that are in the map.
- First posted on arXiv 1 Oct 2021.

## Builds on (most cited first, 26 of 26)

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Dropout as a Bayesian Approximation: Representing Model Uncertainty in Deep Learning](https://mlsanity.com/paper/1506.02142/) (2015)
- [Delving Deep into Rectifiers: Surpassing Human-Level Performance on ImageNet Classification](https://mlsanity.com/paper/1502.01852/) (2015)
- [Variational Dropout and the Local Reparameterization Trick](https://mlsanity.com/paper/1506.02557/) (2015)
- [Pyro: Deep Universal Probabilistic Programming](https://mlsanity.com/paper/1810.09538/) (2018)
- [Variational Continual Learning](https://mlsanity.com/paper/1710.10628/) (2017)
- [Black Box Variational Inference](https://mlsanity.com/paper/1401.0118/) (2013)
- [Probabilistic Backpropagation for Scalable Learning of Bayesian Neural Networks](https://mlsanity.com/paper/1502.05336/) (2015)
- [Multiplicative Normalizing Flows for Variational Bayesian Neural Networks](https://mlsanity.com/paper/1703.01961/) (2017)
- [Flipout: Efficient Pseudo-Independent Weight Perturbations on Mini-Batches](https://mlsanity.com/paper/1803.04386/) (2018)
- [Bayesian graph convolutional neural networks for semi-supervised classification](https://mlsanity.com/paper/1811.11103/) (2018)
- [Structured and Efficient Variational Deep Learning with Matrix Gaussian Posteriors](https://mlsanity.com/paper/1603.04733/) (2016)
- [Bayesian Graph Neural Networks with Adaptive Connection Sampling](https://mlsanity.com/paper/2006.04064/) (2020)
- [Bayesian Layers: A Module for Neural Network Uncertainty](https://mlsanity.com/paper/1812.03973/) (2018)
- [Automated Variational Inference in Probabilistic Programming](https://mlsanity.com/paper/1301.1299/) (2013)
- [Learning Discrete Weights Using the Local Reparameterization Trick](https://mlsanity.com/paper/1710.07739/) (2017)
- [Overpruning in Variational Bayesian Neural Networks](https://mlsanity.com/paper/1801.06230/) (2018)
- [Uncertainty Quantification 360: A Holistic Toolkit for Quantifying and Communicating the Uncertainty of AI](https://mlsanity.com/paper/2106.01410/) (2021)
- [BNNpriors: A library for Bayesian neural network inference with different prior distributions](https://mlsanity.com/paper/2105.06964/) (2021)
- [Simple, Distributed, and Accelerated Probabilistic Programming](https://mlsanity.com/paper/1811.02091/) (2018)
- [Sparse Uncertainty Representation in Deep Learning with Inducing Weights](https://mlsanity.com/paper/2105.14594/) (2021)
- [Improving and Understanding Variational Continual Learning](https://mlsanity.com/paper/1905.02099/) (2019)
- [Learning from the Past: Continual Meta-Learning via Bayesian Graph Modeling](https://mlsanity.com/paper/1911.04695/) (2019)
- [Bayesian Graph Neural Networks for Molecular Property Prediction](https://mlsanity.com/paper/2012.02089/) (2020)
- [Probabilistic Meta-Representations Of Neural Networks](https://mlsanity.com/paper/1810.00555/) (2018)

## Built on by (most cited first, 6 of 6)

- [Position: Bayesian Deep Learning is Needed in the Age of Large-Scale AI](https://mlsanity.com/paper/2402.00809/) (2024)
- [Fortuna: A Library for Uncertainty Quantification in Deep Learning](https://mlsanity.com/paper/2302.04019/) (2023)
- [A Benchmark on Uncertainty Quantification for Deep Learning Prognostics](https://mlsanity.com/paper/2302.04730/) (2023)
- [Black-box Coreset Variational Inference](https://mlsanity.com/paper/2211.02377/) (2022)
- [Density Regression and Uncertainty Quantification with Bayesian Deep Noise Neural Networks](https://mlsanity.com/paper/2206.05643/) (2022)
- [Neetyabhas: A Framework for Uncertainty-Aware Public Policy Optimization in Rational Agent-Based Models](https://mlsanity.com/paper/2606.04562/) (2026)

## Abstract (arXiv)

We introduce TyXe, a Bayesian neural network library built on top of Pytorch and Pyro. Our leading design principle is to cleanly separate architecture, prior, inference and likelihood specification, allowing for a flexible workflow where users can quickly iterate over combinations of these components. In contrast to existing packages TyXe does not implement any layer classes, and instead relies on architectures defined in generic Pytorch code. TyXe then provides modular choices for canonical priors, variational guides, inference techniques, and layer selections for a Bayesian treatment of the specified architecture. Sampling tricks for variance reduction, such as local reparameterization or flipout, are implemented as effect handlers, which can be applied independently of other specifications. We showcase the ease of use of TyXe to explore Bayesian versions of popular models from various libraries: toy regression with a pure Pytorch neural network; large-scale image classification with torchvision ResNets; graph neural networks based on DGL; and Neural Radiance Fields built on top of Pytorch3D. Finally, we provide convenient abstractions for variational continual learning. In all cases the change from a deterministic to a Bayesian neural network comes with minimal modifications to existing code, offering a broad range of researchers and practitioners alike practical access to uncertainty estimation techniques. The library is available at https://github.com/TyXe-BDL/TyXe.