text: # https://medium.com/tensorflow/regression-with-probabilistic-layers-in-tensorflow-probability-e46ff5d37baf
  
  
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  from pprint import pprint
  import matplotlib.pyplot as plt
  import numpy as np
  import seaborn as sns
  
  import tensorflow as tf
  from tensorflow.python import tf2
  if not tf2.enabled():
    import tensorflow.compat.v2 as tf
    tf.enable_v2_behavior()
    assert tf2.enabled()
  
  import tensorflow_probability as tfp
  
  sns.reset_defaults()
  #sns.set_style('whitegrid')
  #sns.set_context('talk')
  sns.set_context(context='talk',font_scale=0.7)
  
  
  tfd = tfp.distributions
  
  negloglik = lambda y, rv_y: -rv_y.log_prob(y)
  
  
  w0 = 0.125
  b0 = 5.
  x_range = [-20, 60]
  
  def load_dataset(n=150, n_tst=150):
    np.random.seed(43)
    def s(x):
      g = (x - x_range[0]) / (x_range[1] - x_range[0])
      return 3 * (0.25 + g**2.)
    x = (x_range[1] - x_range[0]) * np.random.rand(n) + x_range[0]
    eps = np.random.randn(n) * s(x)
    y = (w0 * x * (1. + np.sin(x)) + b0) + eps
    x = x[..., np.newaxis]
    x_tst = np.linspace(*x_range, num=n_tst).astype(np.float32)
    x_tst = x_tst[..., np.newaxis]
    return y, x, x_tst
  
  y, x, x_tst = load_dataset()
  
  # Build model.
  model = tf.keras.Sequential([
    tf.keras.layers.Dense(1),
    tfp.layers.DistributionLambda(lambda t: tfd.Normal(loc=t, scale=1)),
  ])
  
  # Do inference.
  model.compile(optimizer=tf.optimizers.Adam(learning_rate=0.01), loss=negloglik)
  model.fit(x, y, epochs=1000, verbose=False);
  
  # Profit.
  [print(np.squeeze(w.numpy())) for w in model.weights];
  yhat = model(x_tst)
  assert isinstance(yhat, tfd.Distribution)
  
  