def __init__(
self,
shape: Union[int, List[int]],
weight_posterior: Type[BaseDistribution] = Deterministic,
bias_posterior: Type[BaseDistribution] = Deterministic,
weight_prior: BaseDistribution = Normal(0, 1),
bias_prior: BaseDistribution = Normal(0, 1),
weight_initializer: Dict[str, Callable] = {"loc": xavier},
bias_initializer: Dict[str, Callable] = {"loc": xavier},
name="BatchNormalization",
):
# Create the parameters
self.weight = Parameter(
shape=shape,
posterior=weight_posterior,
prior=weight_prior,
initializer=weight_initializer,
name=name + "_weight",
)
self.bias = Parameter(
shape=shape,
posterior=bias_posterior,
prior=bias_prior,
initializer=bias_initializer,
name=name + "_bias",
)
def __call__(self, x):
x = to_tensor(x)
if self.heteroscedastic:
p = x @ self.weights() + self.bias()
m_preds = p[..., :, :self.d_o]
s_preds = O.exp(p[..., :, self.d_o:])
return Normal(m_preds, s_preds)
else:
return Normal(x @ self.weights() + self.bias(), self.std())
def __call__(self, x):
x = to_tensor(x)
if self.heteroscedastic:
p = self.network(x)
m_preds = p[..., :, : self.d_out]
s_preds = O.exp(p[..., :, self.d_out :])
return Normal(m_preds, s_preds)
else:
return Normal(self.network(x), self.std())
def test_r_squared():
"""Tests probflow.utils.metrics.r_squared"""
# Predictive dist
preds = tf.constant([0, 1, 2, 2, 2], dtype=tf.float32)
pred_dist = Normal(preds, 1)
y_true = np.array([0, 1, 2, 3, 4]).astype("float32")
# Compare metric
assert is_close(metrics.r_squared(y_true, pred_dist.mean()), 0.5)
def __call__(self, x):
x = to_tensor(x)
if self.heteroscedastic:
p = self.network(x)
Nd = int(p.shape[-1] / 2)
m_preds = p[..., :, 0:Nd]
s_preds = O.exp(p[..., :, Nd:2 * Nd])
return Normal(m_preds, s_preds)
else:
return Normal(self.network(x), self.std())
def test_sum_squared_error():
"""Tests probflow.utils.metrics.sum_squared_error"""
# Predictive dist
preds = tf.constant([0, 1, 2, 0, 0, 0], dtype=tf.float32)
pred_dist = Normal(preds, 1)
y_true = np.array([0, 0, 0, 0, 1, 2]).astype("float32")
# Compare metric
assert is_close(metrics.sum_squared_error(y_true, pred_dist.mean()), 10.0)
def __init__(self,
shape=1,
posterior=Normal,
prior=Normal(0, 1),
transform=None,
initializer={
'loc': xavier,
'scale': scale_xavier
},
var_transform={
'loc': None,
'scale': O.softplus
},
min: float = 0.0,
max: float = 1.0,
name='BoundedParameter'):
# Check bounds
if min > max:
raise ValueError('min is larger than max')
# Create the transform based on the bounds
if transform is None:
transform = lambda x: min + (max - min) * O.sigmoid(x)
# Create the parameter
super().__init__(shape=shape,
posterior=posterior,
prior=prior,
transform=transform,
initializer=initializer,
var_transform=var_transform,
name=name)
def __init__(
self,
k: Union[int, List[int]],
d: Union[int, List[int]],
posterior: Type[BaseDistribution] = Deterministic,
prior: BaseDistribution = Normal(0, 1),
initializer: Dict[str, Callable] = {"loc": xavier},
name: str = "Embedding",
):
# Convert to list if not already
if isinstance(k, int):
k = [k]
if isinstance(d, int):
d = [d]
# Check values
if len(k) != len(d):
raise ValueError("d and k must be the same length")
if any(e < 1 for e in k):
raise ValueError("k must be >0")
if any(e < 1 for e in d):
raise ValueError("d must be >0")
# Create the parameters
self.embeddings = [
Parameter(
shape=[k[i], d[i]],
posterior=posterior,
prior=prior,
initializer=initializer,
name=name + "_" + str(i),
) for i in range(len(d))
]
def __call__(self, x):
w = self.weight()
b = self.bias()
s = self.std()
m = x * w + b
# check shapes are as expected
if self._is_training:
assert x.ndim == 2
assert x.shape[0] == 1
assert x.shape[1] == 50
assert w.shape[0] == 5
assert w.shape[1] == 1
assert b.shape[0] == 5
assert b.shape[1] == 1
assert s.shape[0] == 5
assert s.shape[1] == 1
assert m.shape[0] == 5
assert m.shape[1] == 50
else: # predicting
assert x.ndim == 1
assert x.shape[0] == 11
assert w.shape[0] == 1
assert b.shape[0] == 1
assert s.shape[0] == 1
assert m.shape[0] == 11
return Normal(m, s)
def __init__(
self,
shape: Union[int, List[int]],
center_by: str = "all",
name="CenteredParameter",
):
# Get a list representing the shape
if isinstance(shape, int):
shape = [shape]
if len(shape) == 1:
shape += [1]
if len(shape) > 2:
raise ValueError(
"Only vector and matrix CenteredParameters are supported")
# Get the untransformed shape of the parameters
if center_by == "row":
K = shape[1]
raw_shape = [K - 1, shape[0]]
elif center_by == "column":
K = shape[0]
raw_shape = [K - 1, shape[1]]
else:
K = shape[0] * shape[1]
raw_shape = [K - 1, 1]
# Prior on the untransformed parameters
scale = float(1.0 / np.sqrt(1 - 1.0 / K))
prior = Normal(0, scale)
# Precompute matrix by which we'll multiply the untransformed params
A = np.eye(K)
A[-1, :] = -1.0
A[-1, -1] = 0.0
Q, _ = np.linalg.qr(A)
self._A_qr = to_default_dtype(to_tensor(Q[:, :-1]))
# Transform function
def A_qr_transform(u):
if center_by == "row":
return O.transpose(self._A_qr @ u)
elif center_by == "all" and shape[1] > 1:
new_shape = list(u.shape) # to handle samples / n_mc > 1
new_shape[-1] = shape[-1]
new_shape[-2] = shape[-2]
return O.reshape(self._A_qr @ u, new_shape)
else:
return self._A_qr @ u
super().__init__(
shape=raw_shape,
posterior=Normal,
prior=prior,
transform=A_qr_transform,
name=name,
)
def test_HiddenMarkovModel():
"""Tests hidden Markov model distribution"""
# Create the distribution (3 states)
initial = tf.random.normal([3])
transition = tf.random.normal([3, 3])
observation = Normal(tf.random.normal([3]), tf.exp(tf.random.normal([3])))
steps = 5
dist = HiddenMarkovModel(initial, transition, observation, steps)
# Should fail w incorrect args
with pytest.raises(TypeError):
HiddenMarkovModel("lala", transition, observation, steps)
with pytest.raises(TypeError):
HiddenMarkovModel(initial, "lala", observation, steps)
with pytest.raises(TypeError):
HiddenMarkovModel(initial, transition, observation, "lala")
with pytest.raises(ValueError):
HiddenMarkovModel(initial, transition, observation, -1)
# Call should return backend obj
assert isinstance(dist(), tfd.HiddenMarkovModel)
# Test sampling
samples = dist.sample()
assert isinstance(samples, tf.Tensor)
assert samples.ndim == 1
assert samples.shape[0] == 5
samples = dist.sample(10)
assert isinstance(samples, tf.Tensor)
assert samples.ndim == 2
assert samples.shape[0] == 10
assert samples.shape[1] == 5
# Test methods
probs = dist.prob([-1.0, 1.0, 0.0, 0.0, 0.0])
assert probs.ndim == 0
probs = dist.prob(np.random.randn(7, 5))
assert probs.ndim == 1
assert probs.shape[0] == 7
# Should also work w/ a backend distribution
observation = tfd.Normal(tf.random.normal([3]),
tf.exp(tf.random.normal([3])))
dist = HiddenMarkovModel(initial, transition, observation, steps)
# Test sampling
samples = dist.sample()
assert isinstance(samples, tf.Tensor)
assert samples.ndim == 1
assert samples.shape[0] == 5
samples = dist.sample(10)
assert isinstance(samples, tf.Tensor)
assert samples.ndim == 2
assert samples.shape[0] == 10
assert samples.shape[1] == 5
def __init__(
self,
shape: Union[int, List[int]] = 1,
posterior: Type[BaseDistribution] = Normal,
prior: BaseDistribution = Normal(0, 1),
transform: Callable = None,
initializer: Dict[str, Callable] = {
"loc": xavier,
"scale": scale_xavier,
},
var_transform: Dict[str, Callable] = {
"loc": None,
"scale": O.softplus,
},
name: str = "Parameter",
):
# Make shape a list
if isinstance(shape, int):
shape = [shape]
# Check values
if any(e < 1 for e in shape):
raise ValueError("all shapes must be >0")
# Assign attributes
self.shape = shape
self.posterior_fn = posterior
self.prior = prior
self.transform = transform if transform else lambda x: x
self.initializer = initializer
self.name = name
self._static_samples_uuid = None
self.var_transform = {
n: (f if f else lambda x: x) for (n, f) in var_transform.items()
}
# Create variables for the variational distribution
self.untransformed_variables = dict()
for var, init in initializer.items():
# Int or float initializations = start whole array at that value
if isinstance(init, (int, float)):
initial_value = O.full(shape, init)
else: # TODO: should also support numpy arrays + backend tensors
initial_value = init(shape)
# Create the variables
self.untransformed_variables[var] = O.new_variable(initial_value)
def __init__(
self,
shape: Union[int, List[int]] = 1,
posterior: Type[BaseDistribution] = Normal,
prior: BaseDistribution = Normal(0, 1),
transform: Callable = None,
initializer: Dict[str, Callable] = {
"loc": xavier,
"scale": scale_xavier,
},
var_transform: Dict[str, Callable] = {
"loc": None,
"scale": O.softplus,
},
name: str = "Parameter",
):
# Make shape a list
if isinstance(shape, int):
shape = [shape]
# Check values
if any(e < 1 for e in shape):
raise ValueError("all shapes must be >0")
# Assign attributes
self.shape = shape
self.posterior_fn = posterior
self.prior = prior
self.transform = transform if transform else lambda x: x
self.initializer = initializer
self.name = name
self.var_transform = {
n: (f if f else lambda x: x)
for (n, f) in var_transform.items()
}
# Create variables for the variational distribution
self.untransformed_variables = dict()
for var, init in initializer.items():
if get_backend() == "pytorch":
import torch
self.untransformed_variables[var] = torch.nn.Parameter(
init(shape))
else:
import tensorflow as tf
self.untransformed_variables[var] = tf.Variable(init(shape))
def __init__(self,
shape=1,
posterior=Deterministic,
prior=Normal(0, 1),
transform=None,
initializer={'loc': xavier},
var_transform={'loc': None},
name='DeterministicParameter'):
super().__init__(shape=shape,
posterior=posterior,
prior=prior,
transform=transform,
initializer=initializer,
var_transform=var_transform,
name=name)
def __init__(
self,
shape=1,
posterior=Normal,
prior=Normal(0, 1),
transform=O.softplus,
initializer={"loc": xavier, "scale": scale_xavier},
var_transform={"loc": None, "scale": O.softplus},
name="PositiveParameter",
):
super().__init__(
shape=shape,
posterior=posterior,
prior=prior,
transform=transform,
initializer=initializer,
var_transform=var_transform,
name=name,
)
def __init__(self,
k: int,
d: int,
posterior: Type[BaseDistribution] = Deterministic,
prior: BaseDistribution = Normal(0, 1),
initializer: Dict[str, Callable] = {'loc': xavier},
name: str = 'Embeddings'):
# Check types
if k < 1:
raise ValueError('k must be >0')
if d < 1:
raise ValueError('d must be >0')
# Create the parameters
self.embeddings = Parameter(shape=[k, d],
posterior=posterior,
prior=prior,
initializer=initializer,
name=name)
def __init__(self,
shape=1,
posterior=Normal,
prior=Normal(0, 1),
transform=O.softplus,
initializer={
'loc': xavier,
'scale': scale_xavier
},
var_transform={
'loc': None,
'scale': O.softplus
},
name='PositiveParameter'):
super().__init__(shape=shape,
posterior=posterior,
prior=prior,
transform=transform,
initializer=initializer,
var_transform=var_transform,
name=name)
def __call__(self, x):
return Normal(self.module(x), self.std())
def __call__(self):
return Normal(self.mean(), self.std())
def __call__(self, x):
return Normal(self.net(x), 1.0)
def __call__(self, x):
return Normal(self.network(x), self.std())
def __call__(self, x):
w = self.weight()
b = self.bias()
s = self.std()
m = x @ w + b
return Normal(m, s)
def __call__(self, x):
x = torch.tensor(x)
return Normal(x @ self.weight() + self.bias(), self.std())
def __call__(self, x):
return Normal(x * self.weight() + self.bias(), 0.1)
def __call__(self, x):
x = to_tensor(x)
return Normal(self.net(x), self.std())
def __call__(self, x):
x = x[self.cols].values
x = to_tensor(x)
return Normal(x @ self.weight() + self.bias(), self.std())
def __call__(self, x):
return Normal(x @ self.weight() + self.bias(), self.std())
def __call__(self, x):
x = torch.tensor(x)
return Normal(self.module(x), self.std())
def __call__(self, x):
return Normal(x, 1.0)
def __call__(self, x):
# can't check shapes b/c tracing it ignores this code
# so just check that it works
return Normal(x * self.weight() + self.bias(), self.std())
