Abstract
We introduce TyXe, a Bayesian neural network library built on top of Pytorchand Pyro. Our leading design principle is to cleanly separate architecture,prior, inference and likelihood specification, allowing for a flexible workflowwhere users can quickly iterate over combinations of these components. Incontrast to existing packages TyXe does not implement any layer classes, andinstead relies on architectures defined in generic Pytorch code. TyXe thenprovides modular choices for canonical priors, variational guides, inferencetechniques, and layer selections for a Bayesian treatment of the specifiedarchitecture. Sampling tricks for variance reduction, such as localreparameterization or flipout, are implemented as effect handlers, which can beapplied independently of other specifications. We showcase the ease of use ofTyXe to explore Bayesian versions of popular models from various libraries: toyregression with a pure Pytorch neural network; large-scale image classificationwith torchvision ResNets; graph neural networks based on DGL; and NeuralRadiance Fields built on top of Pytorch3D. Finally, we provide convenientabstractions for variational continual learning. In all cases the change from adeterministic to a Bayesian neural network comes with minimal modifications toexisting code, offering a broad range of researchers and practitioners alikepractical access to uncertainty estimation techniques. The library is availableat https://github.com/TyXe-BDL/TyXe.