def test_utils_casting():
"""Just gonna toss this in here..."""
df = pd.DataFrame(np.random.randn(5, 2))
o = casting.to_numpy(df)
assert isinstance(o, np.ndarray)
assert o.ndim == 2
assert o.shape[0] == 5
assert o.shape[1] == 2
df = pd.DataFrame(np.random.randn(5, 2))
o = casting.to_tensor(df)
assert o.ndim == 2
assert o.shape[0] == 5
assert o.shape[1] == 2
df = pd.Series(np.random.randn(5))
o = casting.to_tensor(df)
assert o.ndim == 2
assert o.shape[0] == 5
assert o.shape[1] == 1
x = tf.random.normal([5, 2])
o = casting.to_tensor(x)
assert o.ndim == 2
assert o.shape[0] == 5
assert o.shape[1] == 2
def __call__(self, x):
"""Perform the forward pass"""
x = to_tensor(x)
# Using the Flipout estimator
if (get_flipout() and self.flipout and self.probabilistic
and get_samples() is not None and get_samples() == 1):
# Flipout-estimated weight samples
s = O.rand_rademacher(O.shape(x))
r = O.rand_rademacher([O.shape(x)[0], self.d_out])
norm_samples = O.randn([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 = O.rand_rademacher([O.shape(x)[0], self.d_out])
norm_samples = O.randn([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 __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 __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 gather(vals, inds, axis=0):
"""Gather values by index"""
if get_backend() == "pytorch":
import torch
return torch.index_select(vals, axis, to_tensor(inds))
else:
import tensorflow as tf
return tf.gather(vals, inds, axis=axis)
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 __call__(self, x):
x = to_tensor(x)
return Categorical(O.insert_col_of(self.network(x), 0))
def __call__(self, x):
x = to_tensor(x)
return x @ self.weight() + self.bias()
def __call__(self, x):
x = to_tensor(x)
return Normal(x @ self.weight() + self.bias(), self.std())
def prob(self, y):
"""Compute the probability of some data given this distribution"""
if get_backend() == 'pytorch':
return self().log_prob(to_tensor(y)).exp()
else:
return self().prob(to_tensor(y))
def __call__(self, x):
x = to_tensor(x)
for i in range(len(self.layers)):
x = self.layers[i](x)
x = self.activations[i](x)
return x
def __call__(self, x):
x = to_tensor(x)
return Poisson(O.exp(x @ self.weights() + self.bias()))
def __call__(self, x):
x = to_tensor(x)
return Categorical(O.insert_col_of(x @ self.weights() + self.bias(),
0))
def cdf(self, y):
"""Cumulative probability of some data along this distribution"""
return self().cdf(to_tensor(y))
def log_prob(self, y):
"""Compute the log probability of some data given this distribution"""
return self().log_prob(to_tensor(y))
def __call__(self, x):
x = to_tensor(x)
return Bernoulli(x @ self.weight() + self.bias())
def __call__(self, x):
x = to_tensor(x)
return Normal(self.net(x), self.std())
