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