def logdet(self):
z = self.z0 # sxd
u = self.u_ # d
w = self.w_ # d
b = self.b # .
deriv = self.h.deriv # f'
if not self.batched:
# f'(sxd \dot d + .) * -xd = sxd
phi = deriv(z.dot(w) + b).dimshuffle(0, "x") * w.dimshuffle("x", 0)
# \abs(. + sxd \dot d) = s
det = aet.abs_(1.0 + phi.dot(u))
return aet.log(det)
else:
z = z.swapaxes(0, 1)
b = b.dimshuffle(0, "x")
# z bxsxd
# u bxd
# w bxd
# b bx-x-
# f'(bxsxd \bdot bxd + bx-x-) * bx-xd = bxsxd
phi = deriv(aet.batched_dot(z, w) + b).dimshuffle(
0, 1, "x") * w.dimshuffle(0, "x", 1)
# \abs(. + bxsxd \bdot bxd) = bxs
det = aet.abs_(1.0 + aet.batched_dot(phi, u)) # bxs
return aet.log(det).sum(0) # s
def check_jacobian_det(
transform,
domain,
constructor=at.dscalar,
test=0,
make_comparable=None,
elemwise=False,
rv_var=None,
):
y = constructor("y")
y.tag.test_value = test
if rv_var is None:
rv_var = y
rv_inputs = rv_var.owner.inputs if rv_var.owner else []
x = transform.backward(y, *rv_inputs)
if make_comparable:
x = make_comparable(x)
if not elemwise:
jac = at.log(at.nlinalg.det(jacobian(x, [y])))
else:
jac = at.log(at.abs_(at.diag(jacobian(x, [y]))))
# ljd = log jacobian det
actual_ljd = aesara.function([y], jac)
computed_ljd = aesara.function(
[y], at.as_tensor_variable(transform.log_jac_det(y, *rv_inputs)), on_unused_input="ignore"
)
for yval in domain.vals:
close_to(actual_ljd(yval), computed_ljd(yval), tol)
def test_flow_det(flow_spec):
z0 = at.arange(0, 20).astype("float32")
flow = flow_spec(dim=20, z0=z0.dimshuffle("x", 0))
with aesara.config.change_flags(compute_test_value="off"):
z1 = flow.forward.flatten()
J = at.jacobian(z1, z0)
logJdet = at.log(at.abs_(at.nlinalg.det(J)))
det = flow.logdet[0]
np.testing.assert_allclose(logJdet.eval(), det.eval(), atol=0.0001)
def test_flow_det_local(flow_spec):
z0 = at.arange(0, 12).astype("float32")
spec = flow_spec.cls.get_param_spec_for(d=12)
params = dict()
for k, shp in spec.items():
params[k] = np.random.randn(1, *shp).astype("float32")
flow = flow_spec(dim=12, z0=z0.reshape((1, 1, 12)), **params)
assert flow.batched
with aesara.config.change_flags(compute_test_value="off"):
z1 = flow.forward.flatten()
J = at.jacobian(z1, z0)
logJdet = at.log(at.abs_(at.nlinalg.det(J)))
det = flow.logdet[0]
np.testing.assert_allclose(logJdet.eval(), det.eval(), atol=0.0001)
def log_normal(x, mean, **kwargs):
"""
Calculate logarithm of normal distribution at point `x`
with given `mean` and `std`
Parameters
----------
x: Tensor
point of evaluation
mean: Tensor
mean of normal distribution
kwargs: one of parameters `{sigma, tau, w, rho}`
Notes
-----
There are four variants for density parametrization.
They are:
1) standard deviation - `std`
2) `w`, logarithm of `std` :math:`w = log(std)`
3) `rho` that follows this equation :math:`rho = log(exp(std) - 1)`
4) `tau` that follows this equation :math:`tau = std^{-1}`
----
"""
sigma = kwargs.get("sigma")
w = kwargs.get("w")
rho = kwargs.get("rho")
tau = kwargs.get("tau")
eps = kwargs.get("eps", 0.0)
check = sum(map(lambda a: a is not None, [sigma, w, rho, tau]))
if check > 1:
raise ValueError("more than one required kwarg is passed")
if check == 0:
raise ValueError("none of required kwarg is passed")
if sigma is not None:
std = sigma
elif w is not None:
std = aet.exp(w)
elif rho is not None:
std = rho2sigma(rho)
else:
std = tau**(-1)
std += f(eps)
return f(c) - aet.log(aet.abs_(std)) - (x - mean)**2 / (2.0 * std**2)
def _step(i, pkm1, pkm2, qkm1, qkm2, k1, k2, k3, k4, k5, k6, k7, k8, r):
xk = -(x * k1 * k2) / (k3 * k4)
pk = pkm1 + pkm2 * xk
qk = qkm1 + qkm2 * xk
pkm2 = pkm1
pkm1 = pk
qkm2 = qkm1
qkm1 = qk
xk = (x * k5 * k6) / (k7 * k8)
pk = pkm1 + pkm2 * xk
qk = qkm1 + qkm2 * xk
pkm2 = pkm1
pkm1 = pk
qkm2 = qkm1
qkm1 = qk
old_r = r
r = aet.switch(aet.eq(qk, zero), r, pk / qk)
k1 += one
k2 += k26update
k3 += two
k4 += two
k5 += one
k6 -= k26update
k7 += two
k8 += two
big_cond = aet.gt(aet.abs_(qk) + aet.abs_(pk), BIG)
biginv_cond = aet.or_(aet.lt(aet.abs_(qk), BIGINV),
aet.lt(aet.abs_(pk), BIGINV))
pkm2 = aet.switch(big_cond, pkm2 * BIGINV, pkm2)
pkm1 = aet.switch(big_cond, pkm1 * BIGINV, pkm1)
qkm2 = aet.switch(big_cond, qkm2 * BIGINV, qkm2)
qkm1 = aet.switch(big_cond, qkm1 * BIGINV, qkm1)
pkm2 = aet.switch(biginv_cond, pkm2 * BIG, pkm2)
pkm1 = aet.switch(biginv_cond, pkm1 * BIG, pkm1)
qkm2 = aet.switch(biginv_cond, qkm2 * BIG, qkm2)
qkm1 = aet.switch(biginv_cond, qkm1 * BIG, qkm1)
return (
(pkm1, pkm2, qkm1, qkm2, k1, k2, k3, k4, k5, k6, k7, k8, r),
until(aet.abs_(old_r - r) < (THRESH * aet.abs_(r))),
)
def _step(i, t, s):
t *= (i - b) * value / i
step = t / (a + i)
s += step
return ((t, s), until(aet.abs_(step) < threshold))
def sigma2rho(sigma):
"""
`sigma -> rho` aesara converter
:math:`mu + sigma*e = mu + log(1+exp(rho))*e`"""
return aet.log(aet.exp(aet.abs_(sigma)) - 1.0)
def jacobian_det(self, x):
grad = aet.reshape(gradient(aet.sum(self.backward(x)), [x]), x.shape)
return aet.log(aet.abs_(grad))
def dist(self, X, Xs):
if Xs is None:
Xs = at.transpose(X)
else:
Xs = at.transpose(Xs)
return at.abs_((X - Xs + self.c) % (self.c * 2) - self.c)
def laplace(epsilon, obs_data, sim_data):
"""Laplace kernel."""
return -at.abs_((obs_data - sim_data) / epsilon)
def jacobian_det(self, rv_var, rv_value):
grad = at.reshape(
gradient(at.sum(self.backward(rv_var, rv_value)), [rv_value]), rv_value.shape
)
return at.log(at.abs_(grad))
context_name=ctx_name)()
mode = get_mode("FAST_RUN").including("gpuarray")
f = aesara.function([guard_in],
op(guard_in),
mode=mode,
profile=False)
result.cache[key] = f
return f(inp)
result.cache = dict()
return result
f_gpua_min = f_compute(tt.min)
f_gpua_max = f_compute(tt.max)
f_gpua_absmax = f_compute(lambda x: tt.max(tt.abs_(x)))
class NanGuardMode(Mode):
"""
A Aesara compilation Mode that makes the compiled function automatically
detect NaNs and Infs and detect an error if they occur.
Parameters
----------
nan_is_error : bool
If True, raise an error anytime a NaN is encountered.
inf_is_error : bool
If True, raise an error anytime an Inf is encountered. Note that some
pylearn2 modules currently use np.inf as a default value (e.g.
mlp.max_pool) and these will cause an error if inf_is_error is True.
