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]]