tensorflow_probability_normal_distribution_sampling_and_probabilistic_keras_model.py

1import tensorflow as tf
2import tensorflow_probability as tfp
3
4# Create a normal distribution with mean 0 and standard deviation 1.
5dist = tfp.distributions.Normal(loc=0., scale=1.)
6
7# Sample from the distribution
8samples = dist.sample(5)
9
10# Calculate the log-probability of those samples
11log_prob = dist.log_prob(samples)
12
13print(f"Samples: {samples.numpy()}")
14print(f"Log-probabilities: {log_prob.numpy()}")
15
16# Simple linear model with probabilistic layers (from "Get Started" overview)
17model = tf.keras.Sequential([
18    tf.keras.layers.Dense(1),
19    tfp.layers.DistributionLambda(lambda t: tfp.distributions.Normal(loc=t, scale=1)),
20])
21
22# Use the model to predict (outputs a distribution object)
23x = tf.constant([[1.0], [2.0], [3.0]])
24y_dist = model(x)
25print(f"Mean of predictions: {y_dist.mean().numpy()}")