google/uncertainty-baselines: A Forkable Reference for Uncertainty and Robustness Research

The comparison problem this repository was built to solve

The README states the motivation directly: uncertainty and robustness implementations across GitHub are typically one-off experiments for a specific paper, and many papers ship no code at all. The consequence is that even on standard tasks, projects differ in architecture, hyperparameters, and data preprocessing, which makes proper comparison against a baseline difficult. The repository's stated goal is to be a template for researchers to build on, in three ways: high-quality implementations of standard and state-of-the-art methods on standard tasks, minimal dependencies on other files in the codebase so that a baseline is easily forkable, and prescribed best practices for uncertainty and robustness benchmarking. The intended audience is narrow and identifiable: uncertainty and robustness researchers, and anyone who needs to communicate results to them. If you are shipping a production model and want a calibrated confidence score, this is not the repository's target use case, and the absence of a stable release reflects that.

How the baselines, datasets and models modules divide the work

The repository is organized around three importable pieces plus a scripts directory. The baselines/ directory holds the training scripts, organized by training dataset. The README gives baselines/cifar/deterministic.py as an example, described as a Wide ResNet 28-10 obtaining 96.0% test accuracy on CIFAR-10. The ub.datasets module consists of datasets following the TensorFlow Datasets API, adding minimal logic such as default data preprocessing. The ub.models module consists of models following the tf.keras.Model API. The design intent behind the second goal, minimal dependencies on other files, is that a baseline script should be forkable without relying on other baselines and generic modules. That is a deliberate trade-off: it means shared code is kept thin, and a reader can lift one script out of the tree and still understand it. The cost is that improvements to a shared component do not automatically propagate to every baseline, so a fork can drift from the tree it came from.

Loading data and constructing a model in code

The README gives two ways to build a dataset. The first is the class constructor, with CIFAR-10 holding out 10 percent for validation: ub.datasets.Cifar10Dataset(split='train', validation_percent=0.1), followed by dataset_builder.load(batch_size=FLAGS.batch_size), which returns an iterable of batches. The second is the string-based factory, ub.datasets.get(dataset_name, split=split, **dataset_kwargs), which the README notes is useful for instantiating datasets from command line flags. Model construction follows the Keras convention, for example ub.models.wide_resnet(input_shape=(32, 32, 3), depth=28, width_multiplier=10, num_classes=10, l2=1e-4). Two operational notes appear in the README and are easy to miss. In an IPython or Colab notebook, eager execution may need to be activated with tf.compat.v1.enable_eager_execution(). And for Jax or PyTorch consumers, the README shows two conversion paths: tfds.as_numpy(ds), or iterating the dataset and applying jax.tree.map(lambda y: y._numpy(), batch). The README states that tfds.as_numpy calls tensor.numpy(), which invokes an unnecessary copy compared to tensor._numpy(). That second path reaches into a private attribute, so it is a performance suggestion with a compatibility risk attached, not a documented public API.

Installation, backends and the accelerator flag that decides your batch size

Installation is from source: pip install "git+https://github.com/google/uncertainty-baselines.git#egg=uncertainty_baselines". The README is explicit that installing uncertainty_baselines does not automatically install any backend. For TensorFlow you install tensorflow or tf-nightly, tensorflow-addons or tfa-nightly, and tensorboard or tb-nightly; setup.py lists the extra dependencies one can install. There is no stable release and no official release of the library, and the README states all APIs are subject to change. Launching is done through per-baseline flags. The README's Google Cloud example sets BUCKET, TPU_NAME, DATA_DIR and OUTPUT_DIR, then runs python baselines/cifar/batchensemble.py with --tpu, --data_dir and --output_dir. The constraint the README calls out is that the TPU accelerator type must align with the number of cores for the baseline, controlled by the num_cores flag; BatchEnsemble defaults to num_cores=8, so the TPU must be set up with accelerator_type=v3-8. The third option is to change the flags, for example moving to GPU with --use_gpu=True --num_cores=8 and local paths such as --data_dir=/tmp/tensorflow_datasets and --output_dir=/tmp/model. The README's warning here is the most consequential line in the launch section: changing the number of cores matters a lot, because the total batch size during each training step is often determined by num_cores.

What the reported metrics do and do not cover

The README defines the metrics used across datasets, reported to roughly three significant digits and averaged over 10 runs. The first is the number of parameters in the model that makes predictions after training. The second is test accuracy, defined for a dataset of N input-output pairs where the label takes one of K values. The third is test calibration error, the expected calibration error over the test set, attributed to Naeini et al., 2015. The README text supplied here is truncated mid-sentence inside that third definition, so the full formula and any further metrics are not available to quote. That matters for anyone planning to compare numbers: the averaging protocol (10 runs, three significant digits) is stated, but the material does not specify per-dataset splits beyond the validation_percent mechanism, nor how the reported numbers were produced. Treat the metric definitions as the vocabulary the project wants you to use, and treat any specific number you see in a baseline script's comment as a claim to reproduce rather than a guarantee.

Where this repository is the wrong tool

The clearest limitation is stated by the project itself: there is not yet a stable version, there is no official release, and all APIs are subject to change. If you need to pin a dependency in a production pipeline, this repository does not offer that. The second limitation is hardware. The README says you often need TPUs to reproduce baselines, and the three options it lists are Colab, Google Cloud, and changing the flags. Colab offers free TPUs and is described as the most convenient and budget-friendly option, but the README warns against relying on it long-term because TPU access is not guaranteed and Colab can only go so far for managing multiple long experiments. The third limitation is the core-versus-accelerator coupling. Because total batch size is often determined by num_cores, running a baseline on a different core count changes an effective hyperparameter, and the README says results may be similar but ultimately all bets are off. A reader who reduces cores to fit a smaller machine has silently changed the experiment. The fourth is scope: this is a benchmarking template for research comparison, not a calibrated-prediction service, and nothing in the material describes deployment, serving, or monitoring.

How it differs from probabilistic programming frameworks

A natural alternative for uncertainty work is a probabilistic programming framework such as PyMC or NumPyro, and the difference in approach is structural rather than cosmetic. Those frameworks let you declare a generative model and run inference over its parameters, typically on CPU or GPU, with the model expressed in the framework's own language. uncertainty-baselines does the opposite: it fixes the model architecture and the training procedure in a Keras script, then varies the uncertainty method across separate scripts, one per dataset and method. The unit of reuse is a forkable file, not a probabilistic model specification. That makes uncertainty-baselines better suited to asking whether method A beats method B on CIFAR-10 under a shared metric, and worse suited to expressing a bespoke hierarchical model. The repository also depends on the TensorFlow Datasets API for its data layer, so a Jax or PyTorch workflow consumes the datasets through the conversion path shown in the README rather than through a native data loader.

Maintenance, versioning and licence terms

The repository is not archived and the last push recorded in the supplied material is 2026-09-10, but there are no releases retrieved, which is consistent with the README's statement that there is not yet a stable version. The practical upgrade cost follows from that: because installation is from a git URL, the revision you install is whatever the default branch holds at install time, so two installs on different days can resolve to different code. The README also points to setup.py for the extra dependencies, meaning the dependency surface for TensorFlow, TensorFlow Addons and TensorBoard is defined there rather than in a locked environment file. The licence is Apache-2.0. That is a permissive licence, but it carries notice and attribution conditions, and it includes a patent grant and termination clause; if you fork a baseline and redistribute it, the terms you must satisfy are set by the licence text itself. This is not legal advice, and anyone redistributing a fork or shipping it inside a product should read the LICENSE file and their own counsel's guidance rather than this summary.

Editorial conclusion

Adopt uncertainty-baselines if you are an uncertainty or robustness researcher who wants a forkable, per-method training script and a shared metric definition to compare against. Do not adopt it if you need a stable, versioned dependency: the README states there is not yet a stable version and that all APIs are subject to change. Before committing, verify three things: that the baseline script you intend to fork actually exists under baselines/ for your dataset, that your accelerator matches the script's num_cores default, and that you can accept the Apache-2.0 attribution and notice requirements for any redistributed fork.

Official sources

1. google/uncertainty-baselines on GitHub
2. Issues
3. License: Apache-2.0
4. README