$ show tensorflow/probability
TensorFlow Probability (TFP) is a production-grade library for probabilistic modeling and Bayesian inference. It provides a comprehensive suite of probability distributions, bijectors (normalizing flows), MCMC kernels (HMC, NUTS, SMC, etc.), variational inference, Gaussian processes, structural time series, and Bayesian optimization. The library supports TensorFlow, JAX, and NumPy backends through a rewrite system, enabling seamless switching between hardware accelerators. The architecture is modular: each distribution, bijector, and kernel inherits from well-defined abstract base classes with `parameter_properties` for automatic shape and constraint inference. Key entry points include `tfp.distributions`, `tfp.bijectors`, `tfp.mcmc`, `tfp.vi`, `tfp.sts`, and `tfp.experimental` for cutting-edge methods. The code is heavily tested across all backends, with property-based tests and numerical stability checks. While primarily an application library, it also includes research-grade components in `discussion/` and `spinoffs/` (Inference Gym, AutoBNN, FunMC) that serve as benchmarks and experimental prototypes.
The LLM identified 5 files worth highlighting for reuse:
Provides a `structtuple` factory that creates immutable, structurally-typed tuples with caching, slicing, concatenation, and JAX pytree support. This is a self-contained, reusable utility for defining lightweight data containers with robust shape handling, independent of the rest of TFP.
A utility module that tries to compute shapes statically at graph-construction time and falls back to dynamic TensorFlow ops when necessary. It encapsulates the common pattern of preferring static values for performance while maintaining correctness in dynamic contexts. Easily extractable for any TensorFlow or JAX project.
Defines a generic `Reducer` interface with `initialize`, `one_step`, and `finalize` methods for streaming computation of MCMC diagnostics (e.g., R-hat, expectations). This pattern enables efficient, memory-bounded post-processing of Markov chains and is reusable for any streaming aggregation task beyond MCMC.
The abstract base class for all positive semidefinite kernels, providing a consistent API (`apply`, `matrix`, `tensor`) and parameter property descriptors. It demonstrates clean inheritance, matrix construction, and input validation patterns that can be adapted for other kernel-based learning libraries.
Implements `JointDistributionCoroutineAutoBatched`, which automatically infers batch shapes from a probabilistic program written as a generator. This is a non-trivial pattern for composing complex hierarchical models while handling batch dimensions correctly. The design is self-contained and demonstrates advanced use of TensorFlow's `tf.vectorized_map` and shape inference.