def test_wrong_dims(self): a = tt.matrix() increment = tt.matrix() index = 0 with pytest.raises(TypeError): tt.set_subtensor(a[index], increment) with pytest.raises(TypeError): tt.inc_subtensor(a[index], increment)
def test_simple_2d(self): # Increments or sets part of a tensor by a scalar using full slice and # a partial slice depending on a scalar. a = tt.dmatrix() increment = tt.dscalar() sl1 = slice(None) sl2_end = tt.lscalar() sl2 = slice(sl2_end) for do_set in [False, True]: if do_set: resut = tt.set_subtensor(a[sl1, sl2], increment) else: resut = tt.inc_subtensor(a[sl1, sl2], increment) f = aesara.function([a, increment, sl2_end], resut) val_a = np.ones((5, 5)) val_inc = 2.3 val_sl2_end = 2 result = f(val_a, val_inc, val_sl2_end) expected_result = np.copy(val_a) if do_set: expected_result[:, :val_sl2_end] = val_inc else: expected_result[:, :val_sl2_end] += val_inc utt.assert_allclose(result, expected_result)
def test_incsub_f16(): shp = (3, 3) shared = gpuarray_shared_constructor xval = np.arange(np.prod(shp), dtype="float16").reshape(shp) + 1 yval = np.empty((2, ) + shp[1:], dtype="float16") yval[:] = 2 x = shared(xval, name="x") y = tensor.tensor(dtype="float16", broadcastable=(False, ) * len(shp), name="y") expr = tensor.advanced_inc_subtensor1(x, y, [0, 2]) f = aesara.function([y], expr, mode=mode_with_gpu) assert (sum([ isinstance(node.op, GpuAdvancedIncSubtensor1) for node in f.maker.fgraph.toposort() ]) == 1) rval = f(yval) rep = xval.copy() np.add.at(rep, [[0, 2]], yval) assert np.allclose(rval, rep) expr = tensor.inc_subtensor(x[1:], y) f = aesara.function([y], expr, mode=mode_with_gpu) assert (sum([ isinstance(node.op, GpuIncSubtensor) for node in f.maker.fgraph.toposort() ]) == 1) rval = f(yval) rep = xval.copy() rep[1:] += yval assert np.allclose(rval, rep)
def test_incsub_offset(): # Test for https://github.com/Aesara/Aesara/issues/5670 # Build a GPU variable which value will have an offset (x1) x = gpuarray_shared_constructor(np.zeros(5, dtype=aesara.config.floatX)) x1 = x[1:] # Use inc_subtensor on it y = tensor.vector() z = tensor.inc_subtensor(x1[2:], y) # Use updates so that inc_subtensor can happen inplace f = aesara.function([y], z, updates={x: z}, mode=mode_with_gpu) utt.assert_allclose(f([1, 2]), np.array([0, 0, 1, 2], dtype=aesara.config.floatX))
def test_AdvancedIncSubtensor1(x, y, indices): out_aet = aet.set_subtensor(x[indices], y) assert isinstance(out_aet.owner.op, aet_subtensor.AdvancedIncSubtensor1) out_fg = FunctionGraph([], [out_aet]) compare_numba_and_py(out_fg, []) out_aet = aet.inc_subtensor(x[indices], y) assert isinstance(out_aet.owner.op, aet_subtensor.AdvancedIncSubtensor1) out_fg = FunctionGraph([], [out_aet]) compare_numba_and_py(out_fg, []) x_at = x.type() out_aet = aet.set_subtensor(x_at[indices], y, inplace=True) assert isinstance(out_aet.owner.op, aet_subtensor.AdvancedIncSubtensor1) out_fg = FunctionGraph([x_at], [out_aet]) compare_numba_and_py(out_fg, [x.data])
def test_AdvancedIncSubtensor(x, y, indices): out_aet = aet.set_subtensor(x[indices], y) assert isinstance(out_aet.owner.op, aet_subtensor.AdvancedIncSubtensor) out_fg = FunctionGraph([], [out_aet]) compare_numba_and_py(out_fg, []) out_aet = aet.inc_subtensor(x[indices], y) assert isinstance(out_aet.owner.op, aet_subtensor.AdvancedIncSubtensor) out_fg = FunctionGraph([], [out_aet]) compare_numba_and_py(out_fg, []) x_at = x.type() out_aet = aet.set_subtensor(x_at[indices], y) # Inplace isn't really implemented for `AdvancedIncSubtensor`, so we just # hack it on here out_aet.owner.op.inplace = True assert isinstance(out_aet.owner.op, aet_subtensor.AdvancedIncSubtensor) out_fg = FunctionGraph([x_at], [out_aet]) compare_numba_and_py(out_fg, [x.data])
def pos2map(pidx, pgz, prior_result, neib_shape, neib_step): """ Helper function that adds gradient contribution from a single neighborhood position i,j. pidx = Index of position within neighborhood. pgz = Gradient of shape (batch_size*num_channels*neibs) prior_result = Shape (batch_size, num_channnels, rows, cols) neib_shape = Number of rows, cols in a neighborhood. neib_step = Step sizes from image2neibs. """ nrows, ncols = neib_shape rstep, cstep = neib_step batch_size, num_channels, rows, cols = prior_result.shape i = pidx // ncols j = pidx - (i * ncols) # This position does not touch some img pixels in valid mode. result_indices = prior_result[:, :, i:(rows - nrows + i + 1):rstep, j:(cols - ncols + j + 1):cstep, ] newshape = ((batch_size, num_channels) + ((rows - nrows) // rstep + 1, ) + ((cols - ncols) // cstep + 1, )) return tt.inc_subtensor(result_indices, pgz.reshape(newshape))
def forward(self, x): y = aet.zeros(x.shape) y = aet.inc_subtensor(y[..., 0], x[..., 0]) y = aet.inc_subtensor(y[..., 1:], aet.log(x[..., 1:] - x[..., :-1])) return y
def backward(self, y): x = aet.zeros(y.shape) x = aet.inc_subtensor(x[..., 0], y[..., 0]) x = aet.inc_subtensor(x[..., 1:], aet.exp(y[..., 1:])) return aet.cumsum(x, axis=-1)
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 backward(self, value, *inputs): x = at.zeros(value.shape) x = at.inc_subtensor(x[..., 0], value[..., 0]) x = at.inc_subtensor(x[..., 1:], at.exp(value[..., 1:])) return at.cumsum(x, axis=-1)
def just_numeric_args(a, b): return tt.inc_subtensor(a[s], b)
def forward(self, rv_var, rv_value): y = at.zeros(rv_value.shape) y = at.inc_subtensor(y[..., 0], rv_value[..., 0]) y = at.inc_subtensor(y[..., 1:], at.log(rv_value[..., 1:] - rv_value[..., :-1])) return y
def backward(self, rv_var, rv_value): x = at.zeros(rv_value.shape) x = at.inc_subtensor(x[..., 0], rv_value[..., 0]) x = at.inc_subtensor(x[..., 1:], at.exp(rv_value[..., 1:])) return at.cumsum(x, axis=-1)
def forward(self, value, *inputs): y = at.zeros(value.shape) y = at.inc_subtensor(y[..., 0], value[..., 0]) y = at.inc_subtensor(y[..., 1:], at.log(value[..., 1:] - value[..., :-1])) return y
def test_incsubtensor2(self): tv = np.asarray(self.rng.uniform(size=(10,)), aesara.config.floatX) t = aesara.shared(tv) out = tensor.inc_subtensor(t[:4], self.x[:4]) self.check_rop_lop(out, (10,))
def test_incsubtensor1(self): tv = np.asarray(self.rng.uniform(size=(3,)), aesara.config.floatX) t = aesara.shared(tv) out = tensor.inc_subtensor(self.x[:3], t) self.check_rop_lop(out, self.in_shape)