def grad(self, ins, outgrads):
pvals, unis, n = ins
(gz, ) = outgrads
return [
aet.zeros_like(x, dtype=config.floatX)
if x.dtype in discrete_dtypes else aet.zeros_like(x) for x in ins
]
def test_alloc_inputs3():
_W1 = matrix()
_W2 = matrix()
_h0 = vector()
W1 = specify_shape(_W1, (3, 3))
W2 = specify_shape(_W2, (3, 3))
h0 = specify_shape(_h0, (3,))
def lambda_fn(W1, h, W2):
return W1 * dot(h, W2)
o, _ = scan(
lambda_fn,
sequences=at.zeros_like(W1),
outputs_info=h0,
non_sequences=[at.zeros_like(W2)],
n_steps=5,
)
# TODO FIXME: This result depends on unrelated rewrites in the "fast" mode.
f = function([_h0, _W1, _W2], o, mode="FAST_RUN")
scan_node = [x for x in f.maker.fgraph.toposort() if isinstance(x.op, Scan)][0]
assert len(scan_node.op.inner_inputs) == 1
def test_gpujoin_gpualloc():
a = tt.fmatrix("a")
a_val = np.asarray(np.random.rand(4, 5), dtype="float32")
b = tt.fmatrix("b")
b_val = np.asarray(np.random.rand(3, 5), dtype="float32")
f = aesara.function(
[a, b], tt.join(0, tt.zeros_like(a), tt.ones_like(b)) + 4, mode=mode_without_gpu
)
f_gpu = aesara.function(
[a, b], tt.join(0, tt.zeros_like(a), tt.ones_like(b)), mode=mode_with_gpu
)
f_gpu2 = aesara.function(
[a, b], tt.join(0, tt.zeros_like(a), tt.ones_like(b)) + 4, mode=mode_with_gpu
)
assert sum([node.op == tt.alloc for node in f.maker.fgraph.toposort()]) == 2
assert sum([node.op == tt.join_ for node in f.maker.fgraph.toposort()]) == 1
assert (
sum([isinstance(node.op, GpuAlloc) for node in f_gpu.maker.fgraph.toposort()])
== 2
)
assert sum([node.op == gpu_join for node in f_gpu.maker.fgraph.toposort()]) == 1
assert (
sum([isinstance(node.op, GpuAlloc) for node in f_gpu2.maker.fgraph.toposort()])
== 2
)
assert sum([node.op == gpu_join for node in f_gpu2.maker.fgraph.toposort()]) == 1
assert np.allclose(f(a_val, b_val), f_gpu2(a_val, b_val))
def test_alloc_inputs2():
W1 = matrix()
W2 = matrix()
h0 = vector()
def lambda_fn(W1, h, W2):
return W1 * dot(h, W2)
o, _ = scan(
lambda_fn,
sequences=at.zeros_like(W1),
outputs_info=h0,
non_sequences=[at.zeros_like(W2)],
n_steps=5,
)
f = function([h0, W1, W2], o, mode=get_default_mode().including("scan"))
scan_node = [x for x in f.maker.fgraph.toposort() if isinstance(x.op, Scan)][0]
assert (
len(
[
x
for x in scan_node.op.fn.maker.fgraph.toposort()
if isinstance(x.op, Elemwise)
]
)
== 0
)
def moment(rv, size, mu, sigma, init, steps):
grw_moment = at.zeros_like(rv)
grw_moment = at.set_subtensor(grw_moment[..., 0], moment(init))
# Add one dimension to the right, so that mu broadcasts safely along the steps
# dimension
grw_moment = at.set_subtensor(grw_moment[..., 1:], mu[..., None])
return at.cumsum(grw_moment, axis=-1)
def grad(self, inputs, g):
g, g_grad = g
assert str(g_grad) == '<DisconnectedType>'
solution, sens = self(*inputs)
return [
aet.zeros_like(inputs[0]),
aet.sum(g[:, None, :] * sens, (0, -1)),
aesara.gradient.grad_not_implemented(self, 2, inputs[-1])
]
def default_moment(rv, size, *rv_inputs, rv_name=None, has_fallback=False, ndim_supp=0):
if ndim_supp == 0:
return at.zeros(size, dtype=rv.dtype)
elif has_fallback:
return at.zeros_like(rv)
else:
raise TypeError(
"Cannot safely infer the size of a multivariate random variable's moment. "
f"Please provide a moment function when instantiating the {rv_name} "
"random variable."
)
def logp(x, *inputs):
"""Calculate log probability.
Parameters
----------
x: numeric, TensorVariable
Value for which log-probability is calculated.
Returns
-------
TensorVariable
"""
return at.zeros_like(x)
def logp(self, x):
"""Calculate log probability.
Parameters
----------
x: numeric
Value for which log-probability is calculated.
Returns
-------
TensorVariable
"""
return aet.zeros_like(x)
def test_sum_dot(self):
A = matrix("A")
B = matrix("B")
S, _ = scan(
lambda x1, x2, u: u + dot(x1, x2),
sequences=[A.dimshuffle(0, 1, "x"), B.dimshuffle(0, "x", 1)],
outputs_info=[at.zeros_like(A)],
)
f = function([A, B], S.owner.inputs[0][-1])
rng = np.random.default_rng(utt.fetch_seed())
vA = rng.uniform(size=(5, 5)).astype(config.floatX)
vB = rng.uniform(size=(5, 5)).astype(config.floatX)
utt.assert_allclose(f(vA, vB), np.dot(vA.T, vB))
def _logp(
op: Op,
var: TensorVariable,
rvs_to_values: Dict[TensorVariable, TensorVariable],
*inputs: TensorVariable,
**kwargs,
):
"""Create a log-likelihood graph.
This function dispatches on the type of `op`, which should be a subclass
of `RandomVariable`. If you want to implement new log-likelihood graphs
for a `RandomVariable`, register a new function on this dispatcher.
The default assumes that the log-likelihood of a term is a zero.
"""
value_var = rvs_to_values.get(var, var)
return at.zeros_like(value_var)
def grad(self, inputs, g_outputs):
z = at.zeros_like(inputs[0])
gx = inc_diagonal_subtensor(z, inputs[1], inputs[2], g_outputs[0])
return [gx, DisconnectedType()(), DisconnectedType()()]
def isoneutral_diffusion_pre(
maskT,
maskU,
maskV,
maskW,
dxt,
dxu,
dyt,
dyu,
dzt,
dzw,
cost,
cosu,
salt,
temp,
zt,
K_iso,
K_11,
K_22,
K_33,
Ai_ez,
Ai_nz,
Ai_bx,
Ai_by,
):
"""
Isopycnal diffusion for tracer
following functional formulation by Griffies et al
Code adopted from MOM2.1
"""
epsln = 1e-20
iso_slopec = 1e-3
iso_dslope = 1e-3
K_iso_steep = 50.0
tau = 0
dTdx = aet.zeros_like(K_11)
dSdx = aet.zeros_like(K_11)
dTdy = aet.zeros_like(K_11)
dSdy = aet.zeros_like(K_11)
dTdz = aet.zeros_like(K_11)
dSdz = aet.zeros_like(K_11)
"""
drho_dt and drho_ds at centers of T cells
"""
drdT = maskT * get_drhodT(salt[:, :, :, tau], temp[:, :, :, tau], abs(zt))
drdS = maskT * get_drhodS(salt[:, :, :, tau], temp[:, :, :, tau], abs(zt))
"""
gradients at top face of T cells
"""
dTdz = aet.set_subtensor(
dTdz[:, :, :-1],
maskW[:, :, :-1] * (temp[:, :, 1:, tau] - temp[:, :, :-1, tau]) /
dzw[:, :, :-1],
)
dSdz = aet.set_subtensor(
dSdz[:, :, :-1],
maskW[:, :, :-1] * (salt[:, :, 1:, tau] - salt[:, :, :-1, tau]) /
dzw[:, :, :-1],
)
"""
gradients at eastern face of T cells
"""
dTdx = aet.set_subtensor(
dTdx[:-1, :, :],
maskU[:-1, :, :] * (temp[1:, :, :, tau] - temp[:-1, :, :, tau]) /
(dxu[:-1, :, :] * cost[:, :, :]),
)
dSdx = aet.set_subtensor(
dSdx[:-1, :, :],
maskU[:-1, :, :] * (salt[1:, :, :, tau] - salt[:-1, :, :, tau]) /
(dxu[:-1, :, :] * cost[:, :, :]),
)
"""
gradients at northern face of T cells
"""
dTdy = aet.set_subtensor(
dTdy[:, :-1, :],
maskV[:, :-1, :] * (temp[:, 1:, :, tau] - temp[:, :-1, :, tau]) /
dyu[:, :-1, :],
)
dSdy = aet.set_subtensor(
dSdy[:, :-1, :],
maskV[:, :-1, :] * (salt[:, 1:, :, tau] - salt[:, :-1, :, tau]) /
dyu[:, :-1, :],
)
def dm_taper(sx):
"""
tapering function for isopycnal slopes
"""
return 0.5 * (1.0 + aet.tanh((-abs(sx) + iso_slopec) / iso_dslope))
"""
Compute Ai_ez and K11 on center of east face of T cell.
"""
diffloc = aet.zeros_like(K_11)
diffloc = aet.set_subtensor(
diffloc[1:-2, 2:-2, 1:],
0.25 * (K_iso[1:-2, 2:-2, 1:] + K_iso[1:-2, 2:-2, :-1] +
K_iso[2:-1, 2:-2, 1:] + K_iso[2:-1, 2:-2, :-1]),
)
diffloc = aet.set_subtensor(
diffloc[1:-2, 2:-2, 0],
0.5 * (K_iso[1:-2, 2:-2, 0] + K_iso[2:-1, 2:-2, 0]))
sumz = aet.zeros_like(K_11)[1:-2, 2:-2]
for kr in range(2):
ki = 0 if kr == 1 else 1
for ip in range(2):
drodxe = (drdT[1 + ip:-2 + ip, 2:-2, ki:] * dTdx[1:-2, 2:-2, ki:] +
drdS[1 + ip:-2 + ip, 2:-2, ki:] * dSdx[1:-2, 2:-2, ki:])
drodze = (drdT[1 + ip:-2 + ip, 2:-2, ki:] *
dTdz[1 + ip:-2 + ip, 2:-2, :-1 + kr or None] +
drdS[1 + ip:-2 + ip, 2:-2, ki:] *
dSdz[1 + ip:-2 + ip, 2:-2, :-1 + kr or None])
sxe = -drodxe / (aet.minimum(0.0, drodze) - epsln)
taper = dm_taper(sxe)
sumz = aet.inc_subtensor(
sumz[:, :, ki:],
dzw[:, :, :-1 + kr or None] * maskU[1:-2, 2:-2, ki:] *
aet.maximum(K_iso_steep, diffloc[1:-2, 2:-2, ki:] * taper),
)
Ai_ez = aet.set_subtensor(Ai_ez[1:-2, 2:-2, ki:, ip, kr],
taper * sxe * maskU[1:-2, 2:-2, ki:])
K_11 = aet.set_subtensor(K_11[1:-2, 2:-2, :], sumz / (4.0 * dzt[:, :, :]))
"""
Compute Ai_nz and K_22 on center of north face of T cell.
"""
diffloc = aet.set_subtensor(diffloc[...], 0)
diffloc = aet.set_subtensor(
diffloc[2:-2, 1:-2, 1:],
0.25 * (K_iso[2:-2, 1:-2, 1:] + K_iso[2:-2, 1:-2, :-1] +
K_iso[2:-2, 2:-1, 1:] + K_iso[2:-2, 2:-1, :-1]),
)
diffloc = aet.set_subtensor(
diffloc[2:-2, 1:-2, 0],
0.5 * (K_iso[2:-2, 1:-2, 0] + K_iso[2:-2, 2:-1, 0]))
sumz = aet.zeros_like(K_11)[2:-2, 1:-2]
for kr in range(2):
ki = 0 if kr == 1 else 1
for jp in range(2):
drodyn = (drdT[2:-2, 1 + jp:-2 + jp, ki:] * dTdy[2:-2, 1:-2, ki:] +
drdS[2:-2, 1 + jp:-2 + jp, ki:] * dSdy[2:-2, 1:-2, ki:])
drodzn = (drdT[2:-2, 1 + jp:-2 + jp, ki:] *
dTdz[2:-2, 1 + jp:-2 + jp, :-1 + kr or None] +
drdS[2:-2, 1 + jp:-2 + jp, ki:] *
dSdz[2:-2, 1 + jp:-2 + jp, :-1 + kr or None])
syn = -drodyn / (aet.minimum(0.0, drodzn) - epsln)
taper = dm_taper(syn)
sumz = aet.inc_subtensor(
sumz[:, :, ki:],
dzw[:, :, :-1 + kr or None] * maskV[2:-2, 1:-2, ki:] *
aet.maximum(K_iso_steep, diffloc[2:-2, 1:-2, ki:] * taper),
)
Ai_nz = aet.set_subtensor(Ai_nz[2:-2, 1:-2, ki:, jp, kr],
taper * syn * maskV[2:-2, 1:-2, ki:])
K_22 = aet.set_subtensor(K_22[2:-2, 1:-2, :], sumz / (4.0 * dzt[:, :, :]))
"""
compute Ai_bx, Ai_by and K33 on top face of T cell.
"""
sumx = aet.zeros_like(K_11)[2:-2, 2:-2, :-1]
sumy = aet.zeros_like(K_11)[2:-2, 2:-2, :-1]
for kr in range(2):
drodzb = (
drdT[2:-2, 2:-2, kr:-1 + kr or None] * dTdz[2:-2, 2:-2, :-1] +
drdS[2:-2, 2:-2, kr:-1 + kr or None] * dSdz[2:-2, 2:-2, :-1])
# eastward slopes at the top of T cells
for ip in range(2):
drodxb = (drdT[2:-2, 2:-2, kr:-1 + kr or None] *
dTdx[1 + ip:-3 + ip, 2:-2, kr:-1 + kr or None] +
drdS[2:-2, 2:-2, kr:-1 + kr or None] *
dSdx[1 + ip:-3 + ip, 2:-2, kr:-1 + kr or None])
sxb = -drodxb / (aet.minimum(0.0, drodzb) - epsln)
taper = dm_taper(sxb)
sumx += (dxu[1 + ip:-3 + ip, :, :] * K_iso[2:-2, 2:-2, :-1] *
taper * sxb**2 * maskW[2:-2, 2:-2, :-1])
Ai_bx = aet.set_subtensor(Ai_bx[2:-2, 2:-2, :-1, ip, kr],
taper * sxb * maskW[2:-2, 2:-2, :-1])
# northward slopes at the top of T cells
for jp in range(2):
facty = cosu[:, 1 + jp:-3 + jp] * dyu[:, 1 + jp:-3 + jp]
drodyb = (drdT[2:-2, 2:-2, kr:-1 + kr or None] *
dTdy[2:-2, 1 + jp:-3 + jp, kr:-1 + kr or None] +
drdS[2:-2, 2:-2, kr:-1 + kr or None] *
dSdy[2:-2, 1 + jp:-3 + jp, kr:-1 + kr or None])
syb = -drodyb / (aet.minimum(0.0, drodzb) - epsln)
taper = dm_taper(syb)
sumy += (facty * K_iso[2:-2, 2:-2, :-1] * taper * syb**2 *
maskW[2:-2, 2:-2, :-1])
Ai_by = aet.set_subtensor(Ai_by[2:-2, 2:-2, :-1, jp, kr],
taper * syb * maskW[2:-2, 2:-2, :-1])
K_33 = aet.set_subtensor(
K_33[2:-2, 2:-2, :-1],
sumx / (4 * dxt[2:-2, :, :]) + sumy /
(4 * dyt[:, 2:-2, :] * cost[:, 2:-2, :]),
)
K_33 = aet.set_subtensor(K_33[2:-2, 2:-2, -1], 0.0)
return K_11, K_22, K_33, Ai_ez, Ai_nz, Ai_bx, Ai_by
def logpt(
var: TensorVariable,
rv_values: Optional[Union[TensorVariable, Dict[TensorVariable,
TensorVariable]]] = None,
*,
jacobian: bool = True,
scaling: bool = True,
transformed: bool = True,
cdf: bool = False,
sum: bool = False,
**kwargs,
) -> TensorVariable:
"""Create a measure-space (i.e. log-likelihood) graph for a random variable at a given point.
The input `var` determines which log-likelihood graph is used and
`rv_value` is that graph's input parameter. For example, if `var` is
the output of a ``NormalRV`` ``Op``, then the output is a graph of the
density function for `var` set to the value `rv_value`.
Parameters
==========
var
The `RandomVariable` output that determines the log-likelihood graph.
rv_values
A variable, or ``dict`` of variables, that represents the value of
`var` in its log-likelihood. If no `rv_value` is provided,
``var.tag.value_var`` will be checked and, when available, used.
jacobian
Whether or not to include the Jacobian term.
scaling
A scaling term to apply to the generated log-likelihood graph.
transformed
Apply transforms.
cdf
Return the log cumulative distribution.
sum
Sum the log-likelihood.
"""
if not isinstance(rv_values, Mapping):
rv_values = {var: rv_values} if rv_values is not None else {}
rv_var, rv_value_var = extract_rv_and_value_vars(var)
rv_value = rv_values.get(rv_var, rv_value_var)
if rv_var is not None and rv_value is None:
raise ValueError(
f"No value variable specified or associated with {rv_var}")
if rv_value is not None:
rv_value = at.as_tensor(rv_value)
if rv_var is not None:
# Make sure that the value is compatible with the random variable
rv_value = rv_var.type.filter_variable(
rv_value.astype(rv_var.dtype))
if rv_value_var is None:
rv_value_var = rv_value
if rv_var is None:
if var.owner is not None:
return _logp(
var.owner.op,
var,
rv_values,
*var.owner.inputs,
jacobian=jacobian,
scaling=scaling,
transformed=transformed,
cdf=cdf,
sum=sum,
)
return at.zeros_like(var)
rv_node = rv_var.owner
rng, size, dtype, *dist_params = rv_node.inputs
# Here, we plug the actual random variable into the log-likelihood graph,
# because we want a log-likelihood graph that only contains
# random variables. This is important, because a random variable's
# parameters can contain random variables themselves.
# Ultimately, with a graph containing only random variables and
# "deterministics", we can simply replace all the random variables with
# their value variables and be done.
tmp_rv_values = rv_values.copy()
tmp_rv_values[rv_var] = rv_var
if not cdf:
logp_var = _logp(rv_node.op, rv_var, tmp_rv_values, *dist_params,
**kwargs)
else:
logp_var = _logcdf(rv_node.op, rv_var, tmp_rv_values, *dist_params,
**kwargs)
transform = getattr(rv_value_var.tag, "transform",
None) if rv_value_var else None
if transform and transformed and not cdf and jacobian:
transformed_jacobian = transform.jacobian_det(rv_var, rv_value)
if transformed_jacobian:
if logp_var.ndim > transformed_jacobian.ndim:
logp_var = logp_var.sum(axis=-1)
logp_var += transformed_jacobian
# Replace random variables with their value variables
replacements = rv_values.copy()
replacements.update({rv_var: rv_value, rv_value_var: rv_value})
(logp_var, ), _ = rvs_to_value_vars(
(logp_var, ),
apply_transforms=transformed and not cdf,
initial_replacements=replacements,
)
if sum:
logp_var = at.sum(logp_var)
if scaling:
logp_var *= _get_scaling(getattr(rv_var.tag, "total_size", None),
rv_value.shape, rv_value.ndim)
# Recompute test values for the changes introduced by the replacements
# above.
if config.compute_test_value != "off":
for node in io_toposort(graph_inputs((logp_var, )), (logp_var, )):
compute_test_value(node)
if rv_var.name is not None:
logp_var.name = "__logp_%s" % rv_var.name
return logp_var
