def test_sampling():
"""Tests the Sampling context manager"""
# Defaults before sampling
assert settings.get_backend() == "tensorflow"
assert settings.get_samples() is None
assert settings.get_flipout() is False
# Default should be samples=1 and flipout=False
with settings.Sampling():
assert settings.get_samples() == 1
assert settings.get_flipout() is False
# Should return to defaults after sampling
assert settings.get_backend() == "tensorflow"
assert settings.get_samples() is None
assert settings.get_flipout() is False
# Should be able to set samples and flipout via kwargs
with settings.Sampling(n=100, flipout=True):
assert settings.get_samples() == 100
assert settings.get_flipout() is True
# Again should return to defaults after __exit__
assert settings.get_backend() == "tensorflow"
assert settings.get_samples() is None
assert settings.get_flipout() is False
def __init__(self, distributions, logits=None, probs=None):
# Check input
if logits is None and probs is None:
raise ValueError("must pass either logits or probs")
if probs is not None:
ensure_tensor_like(probs, "probs")
if logits is not None:
ensure_tensor_like(logits, "logits")
# Distributions should be a pf, tf, or pt distribution
if not isinstance(distributions, BaseDistribution):
if get_backend() == "pytorch":
import torch.distributions as tod
if not isinstance(distributions, tod.Distribution):
raise TypeError(
"requires either a ProbFlow or PyTorch distribution")
else:
from tensorflow_probability import distributions as tfd
if not isinstance(distributions, tfd.Distribution):
raise TypeError(
"requires either a ProbFlow or TensorFlow distribution"
)
# Store args
self.distributions = distributions
self.logits = logits
self.probs = probs
def to_default_dtype(x):
if get_backend() == "pytorch":
return x.type(get_datatype())
else:
import tensorflow as tf
return tf.cast(x, get_datatype())
def train_step(self, x_data, y_data):
"""Perform one training step"""
elbo = self._train_fn(x_data, y_data)
if get_backend() == "pytorch":
self._current_elbo += elbo.detach().numpy()
else:
self._current_elbo += elbo.numpy()
def log_cholesky_transform(x):
r"""Perform the log cholesky transform on a vector of values.
This turns a vector of :math:`\frac{N(N+1)}{2}` unconstrained values into a
valid :math:`N \times N` covariance matrix.
References
----------
- Jose C. Pinheiro & Douglas M. Bates. `Unconstrained Parameterizations
for Variance-Covariance Matrices
<https://dx.doi.org/10.1007/BF00140873>`_ *Statistics and Computing*,
1996.
"""
if get_backend() == "pytorch":
import numpy as np
import torch
N = int((np.sqrt(1 + 8 * torch.numel(x)) - 1) / 2)
E = torch.zeros((N, N), dtype=get_datatype())
tril_ix = torch.tril_indices(row=N, col=N, offset=0)
E[..., tril_ix[0], tril_ix[1]] = x
E[..., range(N), range(N)] = torch.exp(torch.diagonal(E))
return E @ torch.transpose(E, -1, -2)
else:
import tensorflow as tf
import tensorflow_probability as tfp
E = tfp.math.fill_triangular(x)
E = tf.linalg.set_diag(E, tf.exp(tf.linalg.tensor_diag_part(E)))
return E @ tf.transpose(E)
def __call__(self):
"""Get the distribution object from the backend"""
if get_backend() == "pytorch":
# import torch.distributions as tod
raise NotImplementedError
else:
import tensorflow as tf
from tensorflow_probability import distributions as tfd
# Convert to tensorflow distributions if probflow distributions
if isinstance(self.distributions, BaseDistribution):
self.distributions = self.distributions()
# Broadcast probs/logits
shape = self.distributions.batch_shape
args = {"logits": None, "probs": None}
if self.logits is not None:
args["logits"] = tf.broadcast_to(self["logits"], shape)
else:
args["probs"] = tf.broadcast_to(self["probs"], shape)
# Return TFP distribution object
return tfd.MixtureSameFamily(
tfd.Categorical(**args), self.distributions
)
def kl_divergence(P, Q):
"""Compute the Kullback–Leibler divergence between two distributions.
Parameters
----------
P : |tfp.Distribution| or |torch.Distribution|
The first distribution
Q : |tfp.Distribution| or |torch.Distribution|
The second distribution
Returns
-------
kld : Tensor
The Kullback–Leibler divergence between P and Q (KL(P || Q))
"""
# Get the backend distribution if needed
if isinstance(P, BaseDistribution):
P = P()
if isinstance(Q, BaseDistribution):
Q = Q()
# Compute KL divergence with the backend
if get_backend() == "pytorch":
import torch
return torch.distributions.kl.kl_divergence(P, Q)
else:
import tensorflow_probability as tfp
return tfp.distributions.kl_divergence(P, Q)
def __call__(self, x):
"""Perform the forward pass"""
# Using the Flipout estimator
if get_flipout():
# With PyTorch
if get_backend() == "pytorch":
raise NotImplementedError
# With Tensorflow
else:
import tensorflow as tf
import tensorflow_probability as tfp
# Flipout-estimated weight samples
s = tfp.python.math.random_rademacher(tf.shape(x))
r = tfp.python.math.random_rademacher([x.shape[0], self.d_out])
norm_samples = tf.random.normal([self.d_in, self.d_out])
w_samples = self.weights.variables["scale"] * norm_samples
w_noise = r * ((x * s) @ w_samples)
w_outputs = x @ self.weights.variables["loc"] + w_noise
# Flipout-estimated bias samples
r = tfp.python.math.random_rademacher([x.shape[0], self.d_out])
norm_samples = tf.random.normal([self.d_out])
b_samples = self.bias.variables["scale"] * norm_samples
b_outputs = self.bias.variables["loc"] + r * b_samples
return w_outputs + b_outputs
# Without Flipout
else:
return x @ self.weights() + self.bias()
def square(val):
"""Power of 2"""
if get_backend() == "pytorch":
return val**2
else:
import tensorflow as tf
return tf.math.square(val)
def copy_tensor(x):
"""Copy a tensor, detaching it from the gradient/backend/etc/etc"""
if get_backend() == "pytorch":
return x.detach().clone()
else:
import tensorflow as tf
return tf.identity(x)
def set_learning_rate(self, lr):
"""Set the learning rate used by this model's optimizer"""
if not isinstance(lr, float):
raise TypeError("lr must be a float")
else:
self._learning_rate = lr
if get_backend() == "pytorch":
for g in self._optimizer.param_groups:
g["lr"] = self._learning_rate
def gather(vals, inds, axis=0):
"""Gather values by index"""
if get_backend() == "pytorch":
import torch
return torch.index_select(vals, axis, inds)
else:
import tensorflow as tf
return tf.gather(vals, inds, axis=axis)
def sigmoid(val):
"""Sigmoid function."""
if get_backend() == "pytorch":
import torch
return torch.nn.Sigmoid()(val)
else:
import tensorflow as tf
return tf.math.sigmoid(val)
def softplus(val):
"""Linear rectification."""
if get_backend() == "pytorch":
import torch
return torch.nn.Softplus()(val)
else:
import tensorflow as tf
return tf.math.softplus(val)
def relu(val):
"""Linear rectification."""
if get_backend() == "pytorch":
import torch
return torch.nn.ReLU()(val)
else:
import tensorflow as tf
return tf.nn.relu(val)
def exp(val):
"""The natural exponent."""
if get_backend() == "pytorch":
import torch
return torch.exp(val)
else:
import tensorflow as tf
return tf.exp(val)
def sqrt(val):
"""The square root."""
if get_backend() == "pytorch":
import torch
return torch.sqrt(val)
else:
import tensorflow as tf
return tf.sqrt(val)
def abs(val):
"""Absolute value"""
if get_backend() == "pytorch":
import torch
return torch.abs(val)
else:
import tensorflow as tf
return tf.math.abs(val)
def cat(vals, axis=0):
"""Concatenate tensors"""
if get_backend() == "pytorch":
import torch
return torch.cat(vals, dim=axis)
else:
import tensorflow as tf
return tf.concat(vals, axis=axis)
def prod(val, axis=-1):
"""The product."""
if get_backend() == "pytorch":
import torch
return torch.prod(val, dim=axis)
else:
import tensorflow as tf
return tf.reduce_prod(val, axis=axis)
def eye(dims):
"""Identity matrix."""
if get_backend() == "pytorch":
import torch
return torch.eye(dims)
else:
import tensorflow as tf
return tf.eye(dims)
def zeros(shape):
"""Tensor full of zeros."""
if get_backend() == "pytorch":
import torch
return torch.zeros(shape)
else:
import tensorflow as tf
return tf.zeros(shape)
def squeeze(val):
"""Remove singleton dimensions"""
if get_backend() == "pytorch":
import torch
return torch.squeeze(val)
else:
import tensorflow as tf
return tf.squeeze(val)
def test_backend():
"""Tests setting and getting the backend"""
# Default should be tensorflow
assert settings.get_backend() == "tensorflow"
# Should be able to change to pytorch and back
settings.set_backend("pytorch")
assert settings.get_backend() == "pytorch"
settings.set_backend("tensorflow")
assert settings.get_backend() == "tensorflow"
# But not anything else
with pytest.raises(ValueError):
settings.set_backend("lalala")
# And it has to be a str
with pytest.raises(TypeError):
settings.set_backend(1)
def new_variable(initial_values):
"""Get a new variable with the current backend, and initialize it"""
if get_backend() == "pytorch":
import torch
return torch.nn.Parameter(initial_values)
else:
import tensorflow as tf
return tf.Variable(initial_values)
def reshape(x, new_shape):
"""Reshape a tensor"""
if get_backend() == "pytorch":
import torch
return torch.reshape(x, tuple(new_shape))
else:
import tensorflow as tf
return tf.reshape(x, new_shape)
def mean(val, axis=-1):
"""The mean."""
if get_backend() == "pytorch":
import torch
return torch.mean(val, dim=axis)
else:
import tensorflow as tf
return tf.reduce_mean(val, axis=axis)
def round(val):
"""Round to the closest integer"""
if get_backend() == "pytorch":
import torch
return torch.round(val)
else:
import tensorflow as tf
return tf.math.round(val)
def __call__(self):
"""Get the distribution object from the backend"""
if get_backend() == "pytorch":
import torch.distributions as tod
return tod.dirichlet.Dirichlet(self["concentration"])
else:
from tensorflow_probability import distributions as tfd
return tfd.Dirichlet(self["concentration"])
def std(val, axis=-1):
"""The uncorrected sample standard deviation."""
if get_backend() == "pytorch":
import torch
return torch.std(val, dim=axis)
else:
import tensorflow as tf
return tf.math.reduce_std(val, axis=axis)
