The intent of this notebook is to help TFP 0.11.0 "come to life" via some small snippets - little demos of things you can achieve with TFP.
|---|---|---|---|
| [View on TensorFlow.org](https://www.tensorflow.org/probability/examples/TFP_Release_Notebook_0_11_0) | [Run in Google Colab](https://colab.research.google.com/github/tensorflow/probability/blob/main/tensorflow_probability/examples/jupyter_notebooks/TFP_Release_Notebook_0_11_0.ipynb) | [View source on GitHub](https://github.com/tensorflow/probability/blob/main/tensorflow_probability/examples/jupyter_notebooks/TFP_Release_Notebook_0_11_0.ipynb) | [Download notebook](https://storage.googleapis.com/tensorflow_docs/probability/tensorflow_probability/examples/jupyter_notebooks/TFP_Release_Notebook_0_11_0.ipynb) |
### Installs \& imports
!pip3 install -U -q tensorflow==2.3.0 tensorflow_probability==0.11.0
import tensorflow as tf
import tensorflow_probability as tfp
assert '0.11' in tfp.__version__, tfp.__version__
assert '2.3' in tf.__version__, tf.__version__
tfd = tfp.distributions
tfb = tfp.bijectors
tfpk = tfp.math.psd_kernels
import matplotlib.pyplot as plt
import numpy as np
import scipy.interpolate
import IPython
import seaborn as sns
def plot_spherical(dist, nsamples):
samples = dist.sample([nsamples], seed=123)
probs = dist.prob(samples)
samples = samples.numpy()
probs = probs.numpy()
# Show the dynamic range of the sample probabilities by re-scaling with min/max.
minp, maxp = probs.min(), probs.max()
ranged = np.float64(
(probs - minp) / np.where(maxp - minp > 0, (maxp - minp), minp)) * .8 + .2
data = list(map(list, samples))
colors = map(plt.get_cmap('plasma'), ranged)
colors = list(map(list, colors))
js = """
var scene = new THREE.Scene();
var camera = new THREE.PerspectiveCamera( 75, 1.0, 0.1, 1000 );
var renderer = new THREE.WebGLRenderer();
renderer.setSize( 500, 500 );
document.body.appendChild( renderer.domElement );
var sphere_geom = new THREE.SphereBufferGeometry(1, 200, 200);
var sphere_mat = new THREE.MeshBasicMaterial(
{'opacity': 0.25, 'color': 0xFFFFFF, 'transparent': true});
var sphere = new THREE.Mesh(sphere_geom, sphere_mat);
sphere.position.set(0, 0, 0);
scene.add(sphere);
var points_data = %s;
var points = [];
for (var i = 0; i < points_data.length; i++) {
points.push(new THREE.Vector3(
points_data[i][0], points_data[i][1], points_data[i][2]));
}
var points_geom = new THREE.Geometry();
points_geom.vertices = points;
var colors_data = %s;
var colors = [];
for (var i = 0; i < colors_data.length; i++) {
colors.push(new THREE.Color(
colors_data[i][0], colors_data[i][1], colors_data[i][2]));
}
points_geom.colors = colors;
var points_mat = new THREE.PointsMaterial({'size': 0.015});
points_mat.vertexColors = THREE.VertexColors;
var points = new THREE.Points(points_geom, points_mat);
scene.add(points);
camera.position.x = 0;
camera.position.y = -2;
camera.position.z = 2;
camera.lookAt(new THREE.Vector3(0, 0, 0));
controls = new THREE.TrackballControls(camera, renderer.domElement);
controls.rotateSpeed = 1.0;
controls.zoomSpeed = 1.2;
controls.panSpeed = 0.8;
controls.noZoom = false;
controls.noPan = false;
controls.staticMoving = false;
controls.dynamicDampingFactor = 0.15;
/*
Keep the camera pointing the same direction it was before. However,
this does not generally have the same actual target as the original camera
lookAt, only the same direction. So when the user rotates the view it
might not have the expected origin of rotation.
*/
look_at = new THREE.Vector3(0, 0, -1);
look_at.applyQuaternion(camera.quaternion).add(camera.position)
controls.target = look_at;
var LookAt = function(x, y, z) {
camera.lookAt(new THREE.Vector3(x, y, z));
controls.target.set(x, y, z);
};
LookAt(0, 0, 0);
var animate = function () {
requestAnimationFrame( animate );
controls.update();
renderer.render( scene, camera );
};
animate();
""" % (repr(data), repr(colors))
IPython.display.display_html("""
Spin me!
""" % js, raw=True)
### Visualizing `tfd.SphericalUniform` in 3 dimensions
nsamples = 2000
d = tfd.SphericalUniform(dimension=3)
plot_spherical(d, nsamples)
### Visualizing `tfd.PowerSpherical` in 3 dimensions
stats = ('loc', 'scale', 'tailweight', 'skewness')
plt.figure(figsize=(12, 4))
for i, d in enumerate([
tfd.JohnsonSU(skewness=3, tailweight=2, loc=0, scale=1),
tfd.JohnsonSU(skewness=-3, tailweight=2, loc=0, scale=1)]):
plt.subplot(121+i)
samples = tf.sort(d.sample(2000, seed=(2, 4)))
plt.hist(samples, bins=100, density=True, label='samples')
plt.plot(samples, d.prob(samples), label='p(x)')
plt.legend()
plt.title(f'tfd.JohnsonSU { {k: v for (k, v) in d.parameters.items() if k in stats} }')
plt.tight_layout();
## Auto-vectorized sampling with `JointDistribution*AutoBatched`
If you've ever struggled getting batch dimensions to line up correctly in your joint distributions, this is for you.
### Ever written code like this?
data = tf.random.stateless_normal([7], seed=(1,2))
@tfd.JointDistributionCoroutine
def model():
root = tfd.JointDistributionCoroutine.Root
scale = yield root(tfd.Gamma(1, 1, name='scale'))
yield tfd.Normal(tf.zeros_like(data), scale)
model.sample()
```
(,
)
```
### Then hit exceptions drawing non-scalar samples?
print(model.sample(1), '\n (^^ silent badness: the second value lacks a leading (1,) shape)')
print('sampling (3,) will break:')
import traceback
try:
model.sample(3)
except ValueError as e:
traceback.print_exc()
```
(, )
(^^ silent badness: the second value lacks a leading (1,) shape)
sampling (3,) will break:
Traceback (most recent call last):
File "/usr/local/lib/python3.6/dist-packages/tensorflow_probability/python/distributions/normal.py", line 243, in _parameter_control_dependencies
self._batch_shape()
File "/usr/local/lib/python3.6/dist-packages/tensorflow_probability/python/distributions/normal.py", line 176, in _batch_shape
return tf.broadcast_static_shape(self.loc.shape, self.scale.shape)
File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/util/dispatch.py", line 201, in wrapper
return target(*args, **kwargs)
File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/ops/array_ops.py", line 554, in broadcast_static_shape
return common_shapes.broadcast_shape(shape_x, shape_y)
File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/framework/common_shapes.py", line 107, in broadcast_shape
% (shape_x, shape_y))
ValueError: Incompatible shapes for broadcasting: (7,) and (3,)
During handling of the above exception, another exception occurred:
Traceback (most recent call last):
File "", line 4, in
model.sample(3)
File "/usr/local/lib/python3.6/dist-packages/tensorflow_probability/python/distributions/distribution.py", line 939, in sample
return self._call_sample_n(sample_shape, seed, name, **kwargs)
File "/usr/local/lib/python3.6/dist-packages/tensorflow_probability/python/distributions/joint_distribution.py", line 532, in _call_sample_n
**kwargs)
File "/usr/local/lib/python3.6/dist-packages/tensorflow_probability/python/internal/distribution_util.py", line 1324, in _fn
return fn(*args, **kwargs)
File "/usr/local/lib/python3.6/dist-packages/tensorflow_probability/python/distributions/joint_distribution.py", line 414, in _sample_n
**kwargs)
File "/usr/local/lib/python3.6/dist-packages/tensorflow_probability/python/distributions/joint_distribution.py", line 451, in _call_flat_sample_distributions
ds, xs = self._flat_sample_distributions(sample_shape, seed, value)
File "/usr/local/lib/python3.6/dist-packages/tensorflow_probability/python/distributions/joint_distribution_coroutine.py", line 278, in _flat_sample_distributions
d = gen.send(next_value)
File "", line 7, in model
yield tfd.Normal(tf.zeros_like(data), scale)
File "", line 2, in __init__
File "/usr/local/lib/python3.6/dist-packages/tensorflow_probability/python/distributions/distribution.py", line 334, in wrapped_init
default_init(self_, *args, **kwargs)
File "/usr/local/lib/python3.6/dist-packages/tensorflow_probability/python/distributions/normal.py", line 148, in __init__
name=name)
File "/usr/local/lib/python3.6/dist-packages/tensorflow_probability/python/distributions/distribution.py", line 552, in __init__
d for d in self._parameter_control_dependencies(is_init=True)
File "/usr/local/lib/python3.6/dist-packages/tensorflow_probability/python/distributions/normal.py", line 248, in _parameter_control_dependencies
self.loc.shape, self.scale.shape))
ValueError: Arguments `loc` and `scale` must have compatible shapes; loc.shape=(7,), scale.shape=(3,).
```
### Now, you can use `tfd.JointDistributionCoroutineAutoBatched` \& friends
You write a "scalar" sampling routine, then draw arbitrary sample shapes.
@tfd.JointDistributionCoroutineAutoBatched
def model_auto():
scale = yield tfd.Gamma(1, 1, name='scale')
yield tfd.Normal(tf.zeros_like(data), scale)
print(model_auto.sample(), '\n (scalar sample)')
print(model_auto.sample(1), '\n ((1,) sample)')
print(model_auto.sample(3), '\n ((3,) sample)')
```
(, )
(scalar sample)
WARNING:tensorflow:Note that RandomUniformInt inside pfor op may not give same output as inside a sequential loop.
WARNING:tensorflow:Note that RandomUniformInt inside pfor op may not give same output as inside a sequential loop.
WARNING:tensorflow:Note that RandomStandardNormal inside pfor op may not give same output as inside a sequential loop.
WARNING:tensorflow:Note that RandomStandardNormal inside pfor op may not give same output as inside a sequential loop.
(, )
((1,) sample)
WARNING:tensorflow:Note that RandomUniformInt inside pfor op may not give same output as inside a sequential loop.
WARNING:tensorflow:Note that RandomUniformInt inside pfor op may not give same output as inside a sequential loop.
WARNING:tensorflow:Note that RandomStandardNormal inside pfor op may not give same output as inside a sequential loop.
WARNING:tensorflow:Note that RandomStandardNormal inside pfor op may not give same output as inside a sequential loop.
(, )
((3,) sample)
```
## Reproducible sampling, even in eager
tf.random.set_seed(123)
print('Stateful sampling:\t\t',
tfd.Normal(0, 1).sample(seed=234),
'then',
tfd.Normal(0, 1).sample(seed=234))
# Why are they different?
# TF samplers are stateful by default (TF maintains is a global PRNG).
# How to get reproducible results in eager mode?
tf.random.set_seed(123)
print('Stateful sampling (set_seed):\t',
tfd.Normal(0, 1).sample(seed=234))
# And what about tf.function?
@tf.function
def sample_normal():
return tfd.Normal(0, 1).sample(seed=234)
tf.random.set_seed(123)
print('tf.function:\t\t', sample_normal(), 'vs', sample_normal())
tf.random.set_seed(123)
print('tf.function (set_seed):\t', sample_normal())
# Using a Tensor seed (or a tuple, which TFP will convert to a tensor) induces
# TFP behavior analagous to the tf.random.stateless_* ops. This sampling
# is fully deterministic.
print('Stateless sampling (tuple):\t\t',
tfd.Normal(0, 1).sample(seed=(1, 23)))
print('Stateless sampling (tensor):\t\t',
tfd.Normal(0, 1).sample(seed=tf.constant([1, 23], dtype=tf.int32)))
# Even in tf.function
@tf.function
def sample_normal():
return tfd.Normal(0, 1).sample(seed=(1, 23))
print('Stateless sampling (tf.function):\t', sample_normal())
# And independent of global seeds.
tf.random.set_seed(321)
print('Stateless sampling (ignores set_seed):\t', sample_normal())
```
Stateful sampling: tf.Tensor(0.54054874, shape=(), dtype=float32) then tf.Tensor(-1.5518123, shape=(), dtype=float32)
Stateful sampling (set_seed): tf.Tensor(0.54054874, shape=(), dtype=float32)
tf.function: tf.Tensor(0.54054874, shape=(), dtype=float32) vs tf.Tensor(-1.5518123, shape=(), dtype=float32)
tf.function (set_seed): tf.Tensor(0.54054874, shape=(), dtype=float32)
Stateless sampling (tuple): tf.Tensor(-0.36107817, shape=(), dtype=float32)
Stateless sampling (tensor): tf.Tensor(-0.36107817, shape=(), dtype=float32)
Stateless sampling (tf.function): tf.Tensor(-0.36107817, shape=(), dtype=float32)
Stateless sampling (ignores set_seed): tf.Tensor(-0.36107817, shape=(), dtype=float32)
```
### Similarly, [`tfp.mcmc.sample_chain`](https://www.tensorflow.org/probability/api_docs/python/tfp/mcmc/sample_chain) now accepts stateless seeds
This is passed to underlying kernels, which must be updated to accept a `seed` argument to `one_step` if they make random proposals. (Built-in kernels are already updated.)
kernel = tfp.mcmc.HamiltonianMonteCarlo(lambda x: -(x - .2)**2,
step_size=1.,
num_leapfrog_steps=2)
print_n_lines = lambda x, n: print('\n'.join(repr(x).split('\n')[:n] + ['...']))
print_n_lines(
tfp.mcmc.sample_chain(
num_results=5,
num_burnin_steps=100,
current_state=tf.zeros([3]),
kernel=kernel,
trace_fn=lambda state, kr: kr,
seed=(1, 2)),
17)
print('And again (reproducibly)...')
print_n_lines(
tfp.mcmc.sample_chain(
num_results=5,
num_burnin_steps=100,
current_state=tf.zeros([3]),
kernel=kernel,
trace_fn=lambda state, kr: kr,
seed=(1, 2)),
17)
```
StatesAndTrace(
all_states=,
trace=MetropolisHastingsKernelResults(
accepted_results=UncalibratedHamiltonianMonteCarloKernelResults(
log_acceptance_correction=,
target_log_prob=,
trace=MetropolisHastingsKernelResults(
accepted_results=UncalibratedHamiltonianMonteCarloKernelResults(
log_acceptance_correction=,
target_log_prob=,
target_log_prob=,
seed=[]
),
is_accepted=,
log_accept_ratio=,
proposed_state=1.0,
proposed_results=UncalibratedRandomWalkResults(
log_acceptance_correction=,
target_log_prob=,
seed=
),
extra=[],
seed=
)
```
## CompositeTensor (experimental)
Save on tracing a [`tf.function`](https://www.tensorflow.org/api_docs/python/tf/function). Return distributions from a [`tf.function`](https://www.tensorflow.org/api_docs/python/tf/function).
@tf.function
def get_mean(d):
return d.mean()
print(get_mean(tfp.experimental.as_composite(tfd.BetaBinomial(100., 2., 7.))))
@tf.function
def returns_dist(logits):
return tfp.experimental.as_composite(tfd.Binomial(10., logits=logits))
print(returns_dist(tf.constant(0.)),
returns_dist(tf.constant(0.)).sample(seed=(2, 3)))
```
tf.Tensor(22.222221, shape=(), dtype=float32)
tfp.distributions.BinomialCT("Binomial", batch_shape=[], event_shape=[], dtype=float32) tf.Tensor(5.0, shape=(), dtype=float32)
```
## Did you know: Batch slicing distributions
def demo_batch_slice():
d = tfd.Normal(
loc=tf.random.stateless_normal([100, 20, 3], seed=(1,2)),
scale=1.)
print(d, '\n (original dist)')
print(d[0], '\n (0th element of the first batch dimension)')
print(d[::2, ..., 1], '\n (every other element of first dim x middle element of final dim)\n')
d = tfd.Dirichlet(
concentration=tf.random.stateless_uniform([100, 20, 3, 7], seed=(1,2)))
print(d, '\n (original dist)')
print(d[:, 0], '\n (0th element of the second dimension)')
print(d[::2, ..., 1], '\n (every other element of first dim x middle element of final dim)')
demo_batch_slice()
```
tfp.distributions.Normal("Normal", batch_shape=[100, 20, 3], event_shape=[], dtype=float32)
(original dist)
tfp.distributions.Normal("Normal", batch_shape=[20, 3], event_shape=[], dtype=float32)
(0th element of the first batch dimension)
tfp.distributions.Normal("Normal", batch_shape=[50, 20], event_shape=[], dtype=float32)
(every other element of first dim x middle element of final dim)
tfp.distributions.Dirichlet("Dirichlet", batch_shape=[100, 20, 3], event_shape=[7], dtype=float32)
(original dist)
tfp.distributions.Dirichlet("Dirichlet", batch_shape=[100, 3], event_shape=[7], dtype=float32)
(0th element of the second dimension)
tfp.distributions.Dirichlet("Dirichlet", batch_shape=[50, 20], event_shape=[7], dtype=float32)
(every other element of first dim x middle element of final dim)
```
## Random novelties: Sampling a sigmoidal Cox (nonhomogenous Poisson) process
We can use the thinning procedure proposed in [Adams et al.](https://icml.cc/Conferences/2009/papers/255.pdf)
Note that a similar thinning approach is relevant to some approaches to nonreversible MCMC (e.g. zig-zag and bouncy particle samplers), where an upper bound on intensity can be used to sample nonhomogenous events.
def sigmoidal_cox_process():
V = 20.
lam0 = 3.
n_homog = tfd.Poisson(V * lam0).sample(seed=(0, 1))
homog_locs = tf.sort(tfd.Uniform(0., V).sample(n_homog, seed=(0, 2)))
homog_intensities = tfd.GaussianProcess(
tfpk.MaternThreeHalves(),
index_points=homog_locs[:, tf.newaxis]).sample(seed=(0, 3))
plt.figure(figsize=(15,10))
plt.title('sampling a non-homogenous poisson process')
plt.hlines(lam0, xmin=0., xmax=V)
sigm_int = tf.math.sigmoid(homog_intensities) * lam0
sp = scipy.interpolate.make_interp_spline(homog_locs, sigm_int)
xs = np.linspace(homog_locs[0], homog_locs[-1], 100)
ys = scipy.interpolate.splev(xs, sp)
plt.plot(xs, ys)
plt.scatter(homog_locs, sigm_int, s=5., c='black')
unif = tfd.Uniform(0, lam0).sample(n_homog)
plt.scatter(homog_locs, tf.where(unif < sigm_int, np.nan, unif), c='r', marker='x')
plt.scatter(homog_locs, tf.where(unif > sigm_int, np.nan, unif), c='g', marker='x')
plt.vlines(homog_locs, 0, -.07, colors='gray', )
plt.text(.8, -.06, 'homogenous', dict(fontsize='medium'), horizontalalignment='right')
nonhomog_locs = tf.gather(homog_locs, tf.where(unif < sigm_int))
plt.vlines(nonhomog_locs, -.1, -.17, colors='black')
plt.text(.8, -.16, 'thinned', dict(fontsize='medium'), horizontalalignment='right')
plt.show()
sigmoidal_cox_process()
## Did you know: Batches of GPs
Gaussian Processes are distributions over functions. Here, we will parameterize by a kernel and a set of index points and showcase some of the batching behavior.
(TFP also supports a GP regression model, where the observations at observed index points determine the posterior distribution over functions on a set of index points, see [`tfd.GaussianProcessRegressionModel`](https://www.tensorflow.org/probability/api_docs/python/tfp/distributions/GaussianProcessRegressionModel).)
from mpl_toolkits import mplot3d
x = np.linspace(-6, 6, 9)
y = np.linspace(-6, 6, 9)
X, Y = np.asarray(np.meshgrid(x, y), dtype=np.float32)
print('We can set up a single GP over many examples of 2d feature')
gp = tfd.GaussianProcess(
tfpk.ExponentiatedQuadratic(length_scale=3., feature_ndims=1),
index_points=tf.reshape(tf.stack([X, Y], axis=-1), (-1, 2)),
jitter=1e-3)
print(gp)
Z = tf.reshape(gp.sample(seed=(0,3)), X.shape)
fig = plt.figure(figsize=(10, 6))
ax = plt.axes(projection='3d')
ax.plot_surface(
X, Y, Z, rstride=1, cstride=1, cmap='coolwarm', edgecolor='none');
ax.set_xlabel('x');ax.set_ylabel('y');
ax.set_title('81 x 2 index points (single GP over all pts)')
plt.show()
print('We can have a batch of independent GPs with their own sets of index points.')
batch_gp = tfd.GaussianProcess(
tfpk.ExponentiatedQuadratic(length_scale=3., feature_ndims=1),
index_points=tf.stack([X, Y], axis=-1),
jitter=1e-3)
print(batch_gp)
Z_batch = batch_gp.sample(seed=(5,7))
fig = plt.figure(figsize=(10, 6))
ax = plt.axes(projection='3d')
ax.plot_surface(
X, Y, Z_batch, rstride=1, cstride=1, cmap='coolwarm', edgecolor='none');
ax.set_xlabel('x');ax.set_ylabel('y');
ax.set_title('9 x 9 x 2 index points (9 indep GPs, one at each y position)');
plt.show()
print('We can also have a corresponding batch of kernels.')
kernel = tfpk.ExponentiatedQuadratic(length_scale=3.,
amplitude=tf.linspace(1., 20., 9),
feature_ndims=1)
print(kernel)
batch_gp_kernel = tfd.GaussianProcess(
kernel,
index_points=tf.stack([X, Y], axis=-1),
jitter=1e-3)
print(batch_gp_kernel)
Z_batch_kernel = batch_gp_kernel.sample(seed=(5,7))
fig = plt.figure(figsize=(10, 6))
ax = plt.axes(projection='3d')
ax.plot_surface(
X, Y, Z_batch_kernel, rstride=1, cstride=1, cmap='coolwarm', edgecolor='none');
ax.set_xlabel('x');ax.set_ylabel('y');
ax.set_title('9 x 9 x 2 index points, batch of kernels across y dim');
```
We can set up a single GP over many examples of 2d feature
tfp.distributions.GaussianProcess("GaussianProcess", batch_shape=[], event_shape=[81], dtype=float32)
```
```
We can have a batch of independent GPs with their own sets of index points.
tfp.distributions.GaussianProcess("GaussianProcess", batch_shape=[9], event_shape=[9], dtype=float32)
```
```
We can also have a corresponding batch of kernels.
tfp.math.psd_kernels.ExponentiatedQuadratic("ExponentiatedQuadratic", batch_shape=(9,), feature_ndims=1, dtype=float32)
tfp.distributions.GaussianProcess("GaussianProcess", batch_shape=[9], event_shape=[9], dtype=float32)
```
