def __get_argsort_indices(self, a, axis):
"""
Calculates indices which can be used to reverse sorting operation of
"a" tensor along "axis".
Returns
-------
1d array if axis is None
list of length len(a.shape) otherwise
"""
# The goal is to get gradient wrt input from gradient
# wrt sort(input, axis)
idx = argsort(a, axis, kind=self.kind, order=self.order)
# rev_idx is the reverse of previous argsort operation
rev_idx = argsort(idx, axis, kind=self.kind, order=self.order)
indices = []
axis_data = switch(ge(axis.data, 0), axis.data, a.ndim + axis.data)
for i in range(a.ndim):
index_val = switch(
eq(i, axis_data),
rev_idx,
self.__get_expanded_dim(a, axis, i),
)
indices.append(index_val)
return indices
def L_op(self, inputs, outputs, gradients):
"""
Cholesky decomposition reverse-mode gradient update.
Symbolic expression for reverse-mode Cholesky gradient taken from [#]_
References
----------
.. [#] I. Murray, "Differentiation of the Cholesky decomposition",
http://arxiv.org/abs/1602.07527
"""
dz = gradients[0]
chol_x = outputs[0]
# Replace the cholesky decomposition with 1 if there are nans
# or solve_upper_triangular will throw a ValueError.
if self.on_error == "nan":
ok = ~atm.any(atm.isnan(chol_x))
chol_x = at.switch(ok, chol_x, 1)
dz = at.switch(ok, dz, 1)
# deal with upper triangular by converting to lower triangular
if not self.lower:
chol_x = chol_x.T
dz = dz.T
def tril_and_halve_diagonal(mtx):
"""Extracts lower triangle of square matrix and halves diagonal."""
return at.tril(mtx) - at.diag(at.diagonal(mtx) / 2.0)
def conjugate_solve_triangular(outer, inner):
"""Computes L^{-T} P L^{-1} for lower-triangular L."""
return solve_upper_triangular(
outer.T, solve_upper_triangular(outer.T, inner.T).T
)
s = conjugate_solve_triangular(
chol_x, tril_and_halve_diagonal(chol_x.T.dot(dz))
)
if self.lower:
grad = at.tril(s + s.T) - at.diag(at.diagonal(s))
else:
grad = at.triu(s + s.T) - at.diag(at.diagonal(s))
if self.on_error == "nan":
return [at.switch(ok, grad, np.nan)]
else:
return [grad]
def grad(self, inputs, cost_grad):
"""
In defining the gradient, the Finite Fourier Transform is viewed as
a complex-differentiable function of a complex variable
"""
a = inputs[0]
n = inputs[1]
axis = inputs[2]
grad = cost_grad[0]
if not isinstance(axis, TensorConstant):
raise NotImplementedError(
f"{self.__class__.__name__}: gradient is currently implemented"
" only for axis being an Aesara constant")
axis = int(axis.data)
# notice that the number of actual elements in wrto is independent of
# possible padding or truncation:
elem = arange(0, shape(a)[axis], 1)
# accounts for padding:
freq = arange(0, n, 1)
outer_res = outer(freq, elem)
pow_outer = exp(((-2 * math.pi * 1j) * outer_res) / (1.0 * n))
res = tensordot(grad, pow_outer, (axis, 0))
# This would be simpler but not implemented by aesara:
# res = switch(lt(n, shape(a)[axis]),
# set_subtensor(res[...,n::], 0, False, False), res)
# Instead we resort to that to account for truncation:
flip_shape = list(np.arange(0, a.ndim)[::-1])
res = res.dimshuffle(flip_shape)
res = switch(
lt(n,
shape(a)[axis]),
set_subtensor(
res[n::, ],
0,
False,
False,
),
res,
)
res = res.dimshuffle(flip_shape)
# insures that gradient shape conforms to input shape:
out_shape = (list(np.arange(0, axis)) + [a.ndim - 1] +
list(np.arange(axis, a.ndim - 1)))
res = res.dimshuffle(*out_shape)
return [res, None, None]
def test_jax_logp():
mu = vector("mu")
mu.tag.test_value = np.r_[0.0, 0.0].astype(config.floatX)
tau = vector("tau")
tau.tag.test_value = np.r_[1.0, 1.0].astype(config.floatX)
sigma = vector("sigma")
sigma.tag.test_value = (1.0 / get_test_value(tau)).astype(config.floatX)
value = vector("value")
value.tag.test_value = np.r_[0.1, -10].astype(config.floatX)
logp = (-tau * (value - mu)**2 + log(tau / np.pi / 2.0)) / 2.0
conditions = [sigma > 0]
alltrue = aet_all([aet_all(1 * val) for val in conditions])
normal_logp = aet.switch(alltrue, logp, -np.inf)
fgraph = FunctionGraph([mu, tau, sigma, value], [normal_logp])
compare_jax_and_py(fgraph, [get_test_value(i) for i in fgraph.inputs])
def scan_checkpoints(
fn,
sequences=None,
outputs_info=None,
non_sequences=None,
name="checkpointscan_fn",
n_steps=None,
save_every_N=10,
padding=True,
):
"""Scan function that uses less memory, but is more restrictive.
In :func:`~aesara.scan`, if you compute the gradient of the output
with respect to the input, you will have to store the intermediate
results at each time step, which can be prohibitively huge. This
function allows to do ``save_every_N`` steps of forward computations
without storing the intermediate results, and to recompute them during
the gradient computation.
Notes
-----
Current assumptions:
* Every sequence has the same length.
* If ``n_steps`` is specified, it has the same value as the length of
any sequence.
* The value of ``save_every_N`` divides the number of steps the scan
will run without remainder.
* Only singly-recurrent and non-recurrent outputs are used.
No multiple recurrences.
* Only the last timestep of any output will ever be used.
Parameters
----------
fn
``fn`` is a function that describes the operations involved in one
step of ``scan``. See the documentation of :func:`~aesara.scan`
for more information.
sequences
``sequences`` is the list of Aesara variables or dictionaries
describing the sequences ``scan`` has to iterate over. All
sequences must be the same length in this version of ``scan``.
outputs_info
``outputs_info`` is the list of Aesara variables or dictionaries
describing the initial state of the outputs computed
recurrently.
non_sequences
``non_sequences`` is the list of arguments that are passed to
``fn`` at each steps. One can opt to exclude variable
used in ``fn`` from this list as long as they are part of the
computational graph, though for clarity we encourage not to do so.
n_steps
``n_steps`` is the number of steps to iterate given as an int
or Aesara scalar (> 0). If any of the input sequences do not have
enough elements, scan will raise an error. If n_steps is not provided,
``scan`` will figure out the amount of steps it should run given its
input sequences.
save_every_N
``save_every_N`` is the number of steps to go without storing
the computations of ``scan`` (ie they will have to be recomputed
during the gradient computation).
padding
If the length of the sequences is not a multiple of ``save_every_N``,
the sequences will be zero padded to make this version of ``scan``
work properly, but will also result in a memory copy. It can be
avoided by setting ``padding`` to False, but you need to make
sure the length of the sequences is a multiple of ``save_every_N``.
Returns
-------
tuple
Tuple of the form ``(outputs, updates)`` as in :func:`~aesara.scan`, but
with a small change: It only contain the output at each
``save_every_N`` step. The time steps that are not returned by
this function will be recomputed during the gradient computation
(if any).
See Also
--------
:func:`~aesara.scan`: Looping in Aesara.
"""
# Standardize the format of input arguments
if sequences is None:
sequences = []
elif not isinstance(sequences, list):
sequences = [sequences]
if not isinstance(outputs_info, list):
outputs_info = [outputs_info]
if non_sequences is None:
non_sequences = []
elif not isinstance(non_sequences, list):
non_sequences = [non_sequences]
# Check that outputs_info has no taps:
for element in outputs_info:
if isinstance(element, dict) and "taps" in element:
raise RuntimeError("scan_checkpoints doesn't work with taps.")
# Determine how many steps the original scan would run
if n_steps is None:
n_steps = sequences[0].shape[0]
# Compute the number of steps of the outer scan
o_n_steps = at.cast(ceil(n_steps / save_every_N), "int64")
# Compute the number of steps of the inner scan
i_n_steps = save_every_N * at.ones((o_n_steps, ), "int64")
mod = n_steps % save_every_N
last_n_steps = at.switch(eq(mod, 0), save_every_N, mod)
i_n_steps = set_subtensor(i_n_steps[-1], last_n_steps)
# Pad the sequences if needed
if padding:
# Since padding could be an empty tensor, Join returns a view of s.
join = Join(view=0)
for i, s in enumerate(sequences):
n = s.shape[0] % save_every_N
z = at.zeros((n, s.shape[1:]), dtype=s.dtype)
sequences[i] = join(0, [s, z])
# Establish the input variables of the outer scan
o_sequences = [
s.reshape(
[s.shape[0] / save_every_N, save_every_N] +
[s.shape[i] for i in range(1, s.ndim)],
s.ndim + 1,
) for s in sequences
]
o_sequences.append(i_n_steps)
new_nitsots = [i for i in outputs_info if i is None]
o_nonsequences = non_sequences
def outer_step(*args):
# Separate the received arguments into their respective (seq, outputs
# from previous iterations, nonseqs) categories
i_sequences = list(args[:len(o_sequences)])
i_prev_outputs = list(args[len(o_sequences):-len(o_nonsequences)])
i_non_sequences = list(args[-len(o_nonsequences):])
i_outputs_infos = i_prev_outputs + [
None,
] * len(new_nitsots)
# Call the user-provided function with the proper arguments
results, updates = scan(
fn=fn,
sequences=i_sequences[:-1],
outputs_info=i_outputs_infos,
non_sequences=i_non_sequences,
name=name + "_inner",
n_steps=i_sequences[-1],
)
if not isinstance(results, list):
results = [results]
# Keep only the last timestep of every output but keep all the updates
if not isinstance(results, list):
return results[-1], updates
else:
return [r[-1] for r in results], updates
results, updates = scan(
fn=outer_step,
sequences=o_sequences,
outputs_info=outputs_info,
non_sequences=o_nonsequences,
name=name + "_outer",
n_steps=o_n_steps,
allow_gc=True,
)
return results, updates
def infer_shape(self, fgraph, node, ishapes):
from aesara.tensor.math import eq, maximum, mul
# inputs[1] can contain at most one value of '-1', meaning the actual
# shape of the output will be automatically computed by reshape, so
# that the total number of elements stays the same.
# TODO: Maybe put that formula here?
# It's not trivial, because we would have to check if the product of
# all the non-minus-one shapes is a divisor of the product of the
# original shapes.
# The following expression leads to cycles in feature_shape,
# because it tries to replace the Shape_i node by the switch
# statement, which depends on Shape_i.
# return [tuple([switch(eq(node.inputs[1][i], -1),
# Shape_i(i)(node.outputs[0]),
# node.inputs[1][i])
# for i in range(self.ndim)]
# )]
# Here, we only simplify if the shape (node.inputs[1]) is a constant,
# ideally it would suffice to check that it is always non-negative.
# If current variable is a scalar and its dimensionality should
# change to self.ndim, then use size 1 for all new dimensions.
if len(ishapes[0]) == 0:
return [(1, ) * self.ndim]
requ = node.inputs[1]
input_size = mul(*ishapes[0])
if isinstance(requ, TensorConstant):
requ = list(requ.data)
requ_part = [ele for ele in requ if ele != -1]
crit = len(requ) - len(requ_part)
if crit == 1 and len(requ_part) > 0:
# If there are both 0 and -1 in requ_size, it is impossible
# to determine a right output, but we can at least prevent
# a division by 0. We do not want to keep a negative
# size here as it could lead to further weird errors
# after other optimizations.
requ_size = mul(*requ_part)
missing = input_size // (1 if requ_size == 0 else requ_size)
for i, ele in enumerate(requ):
if ele == -1:
requ[i] = missing
elif crit == 1: # we reshape to -1
requ = [input_size] if ishapes[0] else [1]
elif crit > 1:
raise ValueError("shape argument to Reshape.perform"
" must have at most one entry equal to -1")
return [requ]
else:
requ = [requ[i] for i in range(self.ndim)]
# since new_dims can have negative value (-1), the
# multiplication of all values should be negated
# to give a positive value.
# To avoid optimization complexity, we avoid checking
# for the case when there are two or more '-1' values.
if self.ndim:
requ_size = -mul(*requ)
# If there are both 0 and -1 in requ_size, it is impossible
# to determine a right output, but we can at least prevent
# a division by 0. We do not want to keep a negative
# size here as it could lead to further weird errors
# after other optimizations.
rest_size = input_size // maximum(requ_size, 1)
return [
tuple([
aet.switch(eq(requ[i], -1), rest_size, requ[i])
for i in range(self.ndim)
])
]
def infer_shape(self, fgraph, node, in_shapes):
temp = node.inputs[0]
M = aet.switch(lt(temp, 0), aet.cast(0, temp.dtype), temp)
return [[M]]
def infer_shape(self, node, in_shapes):
temp = node.inputs[0]
M = basic.switch(basic.lt(temp, 0), basic.cast(0, temp.dtype), temp)
return [[M]]
