def test_nested(): notimpl = NotImplementedOp() ifelseifelseif = IfElseIfElseIf() x1 = scalar("x1") x2 = scalar("x2") c1 = scalar("c1") c2 = scalar("c2") t1 = ifelse(c1, x1, notimpl(x2)) t1.name = "t1" t2 = t1 * 10 t2.name = "t2" t3 = ifelse(c2, t2, x1 + t1) t3.name = "t3" t4 = ifelseifelseif(eq(x1, x2), x1, eq(x1, 5), x2, c2, t3, t3 + 0.5) t4.name = "t4" linker = aesara.link.vm.VMLinker(lazy=False) f = function([c1, c2, x1, x2], t4, mode=Mode(linker=linker, optimizer="fast_run")) with pytest.raises(NotImplementedOpException): f(1, 0, np.array(10, dtype=x1.dtype), 0) linker = aesara.link.vm.VMLinker(lazy=True) f = function([c1, c2, x1, x2], t4, mode=Mode(linker=linker, optimizer="fast_run")) assert f(1, 0, np.array(10, dtype=x1.dtype), 0) == 20.5
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 test_RandomVariable_bcast(): rv = RandomVariable("normal", 0, [0, 0], config.floatX, inplace=True) mu = tensor(config.floatX, [True, False, False]) mu.tag.test_value = np.zeros((1, 2, 3)).astype(config.floatX) sd = tensor(config.floatX, [False, False]) sd.tag.test_value = np.ones((2, 3)).astype(config.floatX) s1 = iscalar() s1.tag.test_value = 1 s2 = iscalar() s2.tag.test_value = 2 s3 = iscalar() s3.tag.test_value = 3 s3 = Assert("testing")(s3, eq(s1, 1)) res = rv(mu, sd, size=(s1, s2, s3)) assert res.broadcastable == (False, ) * 3 size = aet.as_tensor((1, 2, 3), dtype=np.int32).astype(np.int64) res = rv(mu, sd, size=size) assert res.broadcastable == (True, False, False) res = rv(0, 1, size=aet.as_tensor(1, dtype=np.int64)) assert res.broadcastable == (True, )
def test_RandomVariable_bcast_specify_shape(): rv = RandomVariable("normal", 0, [0, 0], config.floatX, inplace=True) s1 = aet.as_tensor(1, dtype=np.int64) s2 = iscalar() s2.tag.test_value = 2 s3 = iscalar() s3.tag.test_value = 3 s3 = Assert("testing")(s3, eq(s1, 1)) size = specify_shape(aet.as_tensor([s1, s3, s2, s2, s1]), (5, )) mu = tensor(config.floatX, [False, False, True]) mu.tag.test_value = np.random.normal(size=(2, 2, 1)).astype(config.floatX) std = tensor(config.floatX, [False, True, True]) std.tag.test_value = np.ones((2, 1, 1)).astype(config.floatX) res = rv(mu, std, size=size) assert res.broadcastable == (True, False, False, False, True)
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 broadcast_shape_iter(arrays, **kwargs): """Compute the shape resulting from broadcasting arrays. Parameters ---------- arrays: Iterable[TensorVariable] or Iterable[Tuple[Variable]] An iterable of tensors, or a tuple of shapes (as tuples), for which the broadcast shape is computed. XXX: Do not call this with a generator/iterator; this function will not make copies! arrays_are_shapes: bool (Optional) Indicates whether or not the `arrays` contains shape tuples. If you use this approach, make sure that the broadcastable dimensions are (scalar) constants with the value `1` or `1` exactly. """ one = aesara.scalar.ScalarConstant(aesara.scalar.int64, 1) arrays_are_shapes = kwargs.pop("arrays_are_shapes", False) if arrays_are_shapes: max_dims = max(len(a) for a in arrays) array_shapes = [(one, ) * (max_dims - len(a)) + tuple(one if getattr(sh, "value", sh) == 1 else sh for sh in a) for a in arrays] else: max_dims = max(a.ndim for a in arrays) array_shapes = [(one, ) * (max_dims - a.ndim) + tuple(one if bcast else sh for sh, bcast in zip(a.shape, a.broadcastable)) for a in arrays] result_dims = [] for dim_shapes in zip(*array_shapes): non_bcast_shapes = [shape for shape in dim_shapes if shape != one] if len(non_bcast_shapes) > 0: # Either there's only one non-broadcastable dimensions--and that's # what determines the dimension size, or there are multiple # non-broadcastable dimensions that must be equal i_dim = non_bcast_shapes.pop() potentially_unequal_dims = [ dim for dim in non_bcast_shapes # TODO FIXME: This is a largely deficient means of comparing graphs # (and especially shapes) if not equal_computations([i_dim], [dim]) ] if potentially_unequal_dims: # In this case, we can't tell whether or not the dimensions are # equal, so we'll need to assert their equality and move the error # handling to evaluation time. assert_dim = Assert("Could not broadcast dimensions") eq_condition = aet_all([ or_(eq(dim, one), eq(i_dim, dim)) for dim in potentially_unequal_dims ]) eq_condition = or_(eq(i_dim, one), eq_condition) result_dims.append(assert_dim(i_dim, eq_condition)) else: result_dims.append(i_dim) else: # Every array was broadcastable in this dimension result_dims.append(one) return tuple(result_dims)
def test_RandomVariable_basics(): str_res = str( RandomVariable( "normal", 0, [0, 0], config.floatX, inplace=True, )) assert str_res == "normal_rv" # `ndims_params` should be a `Sequence` type with raises(TypeError, match="^Parameter ndims_params*"): RandomVariable( "normal", 0, 0, config.floatX, inplace=True, ) # `size` should be a `Sequence` type with raises(TypeError, match="^Parameter size*"): RandomVariable( "normal", 0, [0, 0], config.floatX, inplace=True, )(0, 1, size={1, 2}) # No dtype with raises(TypeError, match="^dtype*"): RandomVariable( "normal", 0, [0, 0], inplace=True, )(0, 1) # Confirm that `inplace` works rv = RandomVariable( "normal", 0, [0, 0], "normal", inplace=True, ) assert rv.inplace assert rv.destroy_map == {0: [3]} # A no-params `RandomVariable` rv = RandomVariable(name="test_rv", ndim_supp=0, ndims_params=()) with raises(TypeError): rv.make_node(rng=1) # `RandomVariable._infer_shape` should handle no parameters rv_shape = rv._infer_shape(aet.constant([]), (), []) assert rv_shape.equals(aet.constant([], dtype="int64")) # Integer-specificed `dtype` dtype_1 = all_dtypes[1] rv_node = rv.make_node(None, None, 1) rv_out = rv_node.outputs[1] rv_out.tag.test_value = 1 assert rv_out.dtype == dtype_1 with raises(NullTypeGradError): grad(rv_out, [rv_node.inputs[0]]) rv = RandomVariable("normal", 0, [0, 0], config.floatX, inplace=True) mu = tensor(config.floatX, [True, False, False]) mu.tag.test_value = np.zeros((1, 2, 3)).astype(config.floatX) sd = tensor(config.floatX, [False, False]) sd.tag.test_value = np.ones((2, 3)).astype(config.floatX) s1 = iscalar() s1.tag.test_value = 1 s2 = iscalar() s2.tag.test_value = 2 s3 = iscalar() s3.tag.test_value = 3 s3 = Assert("testing")(s3, eq(s1, 1)) res = rv.compute_bcast([mu, sd], (s1, s2, s3)) assert res == [False] * 3