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)
