def make_node(self, x, shape): if not isinstance(x, Variable): x = at.as_tensor_variable(x) shape = at.as_tensor_variable(shape, ndim=1) if isinstance(shape, Constant): shape = tuple(shape.data) else: shape = tuple(at.as_tensor_variable(s, ndim=0) for s in shape) if any(s.dtype not in aesara.tensor.type.integer_dtypes for s in shape): raise TypeError("Shape values must be integer types") if len(shape) != x.ndim: raise ValueError( f"Input `x` is {x.ndim}-dimensional and will never match a shape of length {len(shape)}." ) if isinstance(x.type, TensorType) and all( isinstance(s, Number) for s in shape): out_var = TensorType(x.type.dtype, shape)() else: out_var = x.type() in_shape = at.as_tensor_variable(shape, ndim=1) return Apply(self, [x, in_shape], [out_var])
def make_node(self, inp, out_grad, ws, stride=None, pad=None): ctx_name = infer_context_name(inp, out_grad) nd = self.ndim inp = as_gpuarray_variable(inp, ctx_name) assert inp.ndim == nd + 2 out_grad = as_gpuarray_variable(out_grad, ctx_name) assert out_grad.ndim == nd + 2 assert out_grad.ndim == inp.ndim if stride is None: stride = ws if pad is None: pad = (0, ) * nd elif isinstance(pad, (tuple, list)): if max(pad) != 0 and not self.mode == "average_exc_pad": raise ValueError("Padding must be zero for average_exc_pad") ws = as_tensor_variable(ws) stride = as_tensor_variable(stride) pad = as_tensor_variable(pad) assert ws.ndim == stride.ndim and ws.ndim == pad.ndim assert ws.ndim == 1 if ws.dtype not in int_dtypes: raise TypeError("Window shape parameters must be ints.") if stride.dtype not in int_dtypes: raise TypeError("Stride parameters must be ints.") if pad.dtype not in int_dtypes: raise TypeError("Padding parameters must be ints.") ws = aesara.tensor.cast(ws, "int64") stride = aesara.tensor.cast(stride, "int64") pad = aesara.tensor.cast(pad, "int64") return Apply(self, [inp, out_grad, ws, stride, pad], [inp.type()])
def make_node(self, inp, ws, stride=None, pad=None): ctx_name = infer_context_name(inp) inp = as_gpuarray_variable(inp, ctx_name) nd = self.ndim assert inp.ndim == nd + 2 if stride is None: stride = ws if pad is None: pad = (0, ) * nd elif isinstance(pad, (tuple, list)): if max(pad) != 0 and not self.ignore_border: raise ValueError("Padding works only with ignore_border=True") if isinstance(ws, (tuple, list)): if any(pad[i] >= ws[i] for i in range(nd)): raise ValueError("Padding must be smaller than strides") ws = as_tensor_variable(ws) stride = as_tensor_variable(stride) pad = as_tensor_variable(pad) assert ws.ndim == stride.ndim and ws.ndim == pad.ndim assert ws.ndim == 1 if ws.dtype not in int_dtypes: raise TypeError("Window shape parameters must be ints.") if stride.dtype not in int_dtypes: raise TypeError("Stride parameters must be ints.") if pad.dtype not in int_dtypes: raise TypeError("Padding parameters must be ints.") ws = aesara.tensor.cast(ws, "int64") stride = aesara.tensor.cast(stride, "int64") pad = aesara.tensor.cast(pad, "int64") return Apply(self, [inp, ws, stride, pad], [inp.type()])
def make_node(self, x, *shape): from aesara.tensor.basic import get_scalar_constant_value x = at.as_tensor_variable(x) shape = tuple(NoneConst if (s is None or NoneConst.equals(s) ) else at.as_tensor_variable(s, ndim=0) for s in shape) if any(s.dtype not in aesara.tensor.type.integer_dtypes for s in shape if hasattr(s, "dtype")): raise TypeError("Shape values must be integer types") if len(shape) != x.type.ndim: raise ValueError( f"Input `x` is {x.type.ndim}-dimensional and will never match a shape of length {len(shape)}." ) type_shape = [None] * x.ndim for i, (xts, s) in enumerate(zip(x.type.shape, shape)): if xts is not None: type_shape[i] = xts else: try: type_s = get_scalar_constant_value(s) if type_s is not None: type_shape[i] = int(type_s) except NotScalarConstantError: pass out_var = x.type.clone(shape=type_shape)() return Apply(self, [x, *shape], [out_var])
def make_node(self, inp, out, out_grad, ws, stride=None, pad=None): ctx_name = infer_context_name(inp, out, out_grad) nd = self.ndim inp = as_gpuarray_variable(inp, ctx_name) assert inp.ndim == nd + 2 out = as_gpuarray_variable(out, ctx_name) assert out_grad.ndim == nd + 2 out_grad = as_gpuarray_variable(out_grad, ctx_name) assert out.ndim == nd + 2 assert out_grad.ndim == inp.ndim assert inp.ndim == out.ndim if stride is None: stride = ws if pad is None: pad = (0, ) * nd ws = as_tensor_variable(ws) stride = as_tensor_variable(stride) pad = as_tensor_variable(pad) assert ws.ndim == stride.ndim and ws.ndim == pad.ndim assert ws.ndim == 1 if ws.dtype not in int_dtypes: raise TypeError("Window shape parameters must be ints.") if stride.dtype not in int_dtypes: raise TypeError("Stride parameters must be ints.") if pad.dtype not in int_dtypes: raise TypeError("Padding parameters must be ints.") ws = aesara.tensor.cast(ws, "int64") stride = aesara.tensor.cast(stride, "int64") pad = aesara.tensor.cast(pad, "int64") return Apply(self, [inp, out, out_grad, ws, stride, pad], [inp.type()])
def test_jax_compile_ops(): x = DeepCopyOp()(aet.as_tensor_variable(1.1)) x_fg = FunctionGraph([], [x]) compare_jax_and_py(x_fg, []) x_np = np.zeros((20, 1, 1)) x = aet.Rebroadcast((0, False), (1, True), (2, False))(aet.as_tensor_variable(x_np)) x_fg = FunctionGraph([], [x]) compare_jax_and_py(x_fg, []) with config.change_flags(compute_test_value="off"): x = aet.Rebroadcast((0, True), (1, False), (2, False))(aet.as_tensor_variable(x_np)) x_fg = FunctionGraph([], [x]) with pytest.raises(ValueError): compare_jax_and_py(x_fg, []) x = ViewOp()(aet.as_tensor_variable(x_np)) x_fg = FunctionGraph([], [x]) compare_jax_and_py(x_fg, [])
def make_node(self, x, shp): x = aet.as_tensor_variable(x) shp_orig = shp shp = aet.as_tensor_variable(shp, ndim=1) if not (shp.dtype in int_dtypes or (isinstance(shp, TensorConstant) and shp.data.size == 0)): # It raises an error if shp is not of integer type, # except when shp is constant and empty # (in this case, shp.dtype does not matter anymore). raise TypeError("Shape must be integers", shp, shp.dtype) assert shp.ndim == 1 if isinstance(shp, TensorConstant): bcast = [s == 1 for s in shp.data] return Apply(self, [x, shp], [tensor(x.type.dtype, bcast)]) else: bcasts = [False] * self.ndim shp_list = shp_orig if hasattr(shp_orig, "ndim") and shp_orig.ndim == 0: shp_list = [shp_orig] for index in range(self.ndim): y = shp_list[index] y = aet.as_tensor_variable(y) # Try to see if we can infer that y has a constant value of 1. # If so, that dimension should be broadcastable. try: bcasts[index] = (hasattr(y, "get_scalar_constant_value") and y.get_scalar_constant_value() == 1) except NotScalarConstantError: pass return Apply(self, [x, shp], [tensor(x.type.dtype, bcasts)])
def make_node(self, a, val, offset): a = aet.as_tensor_variable(a) val = aet.as_tensor_variable(val) offset = aet.as_tensor_variable(offset) if a.ndim != 2: raise TypeError( "%s: first parameter must have exactly" " two dimensions" % self.__class__.__name__ ) elif val.ndim != 0: raise TypeError( f"{self.__class__.__name__}: second parameter must be a scalar" ) elif offset.ndim != 0: raise TypeError( f"{self.__class__.__name__}: third parameter must be a scalar" ) val = aet.cast(val, dtype=upcast(a.dtype, val.dtype)) if val.dtype != a.dtype: raise TypeError( "%s: type of second parameter must be the same" " as the first's" % self.__class__.__name__ ) elif offset.dtype not in integer_dtypes: raise TypeError( f"{self.__class__.__name__}: type of third parameter must be as integer" " use aesara.tensor.cast( input, 'int32/int64')" ) return Apply(self, [a, val, offset], [a.type()])
def make_node(self, kern, topgrad, shape=None): kern = as_tensor_variable(kern) topgrad = as_tensor_variable(topgrad) kern, topgrad = self.as_common_dtype(kern, topgrad) if kern.type.ndim != 5: raise TypeError("kern must be 5D tensor") if topgrad.type.ndim != 5: raise TypeError("topgrad must be 5D tensor") if shape is None: if self.subsample != (1, 1, 1): raise ValueError("shape must be given if subsample != (1, 1, 1)") height_width_depth = [] else: height_width_depth = [ as_tensor_variable(shape[0]).astype("int64"), as_tensor_variable(shape[1]).astype("int64"), as_tensor_variable(shape[2]).astype("int64"), ] if self.num_groups > 1: broadcastable = [topgrad.type.broadcastable[0], False, False, False, False] else: broadcastable = [ topgrad.type.broadcastable[0], kern.type.broadcastable[1], False, False, False, ] dtype = kern.type.dtype return Apply( self, [kern, topgrad] + height_width_depth, [TensorType(dtype, broadcastable)()], )
def make_node(self, img, topgrad, shape=None): img = as_tensor_variable(img) topgrad = as_tensor_variable(topgrad) img, topgrad = self.as_common_dtype(img, topgrad) if img.type.ndim != 5: raise TypeError("img must be 5D tensor") if topgrad.type.ndim != 5: raise TypeError("topgrad must be 5D tensor") if shape is None: if self.subsample != (1, 1, 1) or self.border_mode == "half": raise ValueError( "shape must be given if subsample != (1, 1, 1)" ' or border_mode == "half"') height_width_depth = [] else: height_width_depth = [ as_tensor_variable(shape[0]).astype("int64"), as_tensor_variable(shape[1]).astype("int64"), as_tensor_variable(shape[2]).astype("int64"), ] broadcastable = [ topgrad.type.broadcastable[1], img.type.broadcastable[1], False, False, False, ] dtype = img.type.dtype return Apply( self, [img, topgrad] + height_width_depth, [TensorType(dtype, broadcastable)()], )
def make_node(self, input, axis=-1): input = as_tensor_variable(input) axis = as_tensor_variable(axis) return Apply( self, [input, axis], [TensorType(dtype="int64", shape=input.type.shape)()], )
def make_node(self, a, *shape): a = aet.as_tensor_variable(a) shape = aet.as_tensor_variable(shape, ndim=1) shape, bcast = aet.alloc_validate_shape(shape) out = type(a.type)(dtype=a.type.dtype, broadcastable=bcast)() return Apply(self, [a] + shape, [out])
def make_node(self, input, axis=-1): input = as_tensor_variable(input) axis = as_tensor_variable(axis) bcast = input.type.broadcastable return Apply( self, [input, axis], [TensorType(dtype="int64", broadcastable=bcast)()], )
def test_jax_shape_ops(): x_np = np.zeros((20, 3)) x = Shape()(aet.as_tensor_variable(x_np)) x_fg = FunctionGraph([], [x]) compare_jax_and_py(x_fg, [], must_be_device_array=False) x = Shape_i(1)(aet.as_tensor_variable(x_np)) x_fg = FunctionGraph([], [x]) compare_jax_and_py(x_fg, [], must_be_device_array=False)
def make_node(self, x, shape): if not isinstance(x, Variable): x = aet.as_tensor_variable(x) shape = aet.as_tensor_variable(shape) if shape.ndim > 1: raise AssertionError() if shape.dtype not in aesara.tensor.type.integer_dtypes: raise AssertionError() if isinstance(shape, TensorConstant) and shape.data.size != x.ndim: raise AssertionError() return Apply(self, [x, shape], [x.type()])
def make_node(self, a, val): a = basic.as_tensor_variable(a) val = basic.as_tensor_variable(val) if a.ndim < 2: raise TypeError("%s: first parameter must have at least" " two dimensions" % self.__class__.__name__) elif val.ndim != 0: raise TypeError("%s: second parameter must be a scalar" % self.__class__.__name__) val = basic.cast(val, dtype=upcast(a.dtype, val.dtype)) if val.dtype != a.dtype: raise TypeError("%s: type of second parameter must be the same as" " the first's" % self.__class__.__name__) return Apply(self, [a, val], [a.type()])
def test_jax_specify_shape(): x_np = np.zeros((20, 3)) x = SpecifyShape()(aet.as_tensor_variable(x_np), (20, 3)) x_fg = FunctionGraph([], [x]) compare_jax_and_py(x_fg, []) with config.change_flags(compute_test_value="off"): x = SpecifyShape()(aet.as_tensor_variable(x_np), (2, 3)) x_fg = FunctionGraph([], [x]) with pytest.raises(AssertionError): compare_jax_and_py(x_fg, [])
def shape(x: Union[np.ndarray, Number, Variable]) -> Variable: """Return the shape of `x`.""" if not isinstance(x, Variable): x = at.as_tensor_variable(x) x_type = x.type if isinstance(x_type, TensorType) and all(s is not None for s in x_type.shape): res = at.as_tensor_variable(x_type.shape, ndim=1, dtype=np.int64) else: res = _shape(x) return res
def make_node(self, x, y, rcond): x = as_tensor_variable(x) y = as_tensor_variable(y) rcond = as_tensor_variable(rcond) return Apply( self, [x, y, rcond], [ matrix(), dvector(), lscalar(), dvector(), ], )
def make_node(self, activations, labels, input_lengths): context_name = infer_context_name(activations) t_activations = as_gpuarray_variable(activations, context_name=context_name) # Ensure activations array is C-contiguous t_activations = gpu_contiguous(t_activations) # Labels and input lengths are always on the CPU t_labels = as_tensor_variable(labels) t_input_lengths = as_tensor_variable(input_lengths) if t_activations.type.dtype != "float32": raise TypeError("activations must use the float32 type.") if t_activations.ndim != 3: raise ValueError("activations must have 3 dimensions.") if t_labels.type.dtype != "int32": raise TypeError("labels must use the int32 type.") if t_labels.ndim != 2: raise ValueError("labels must have 2 dimensions.") if t_input_lengths.type.dtype != "int32": raise TypeError("input_lengths must use the int32 type.") if t_input_lengths.ndim != 1: raise ValueError("input_lengths must have 1 dimension.") costs = GpuArrayType(dtype="float32", broadcastable=(False, ), context_name=context_name)() outputs = [costs] if self.compute_grad: gradients = GpuArrayType( dtype="float32", broadcastable=( False, False, False, ), context_name=context_name, )() outputs += [gradients] return Apply(self, inputs=[t_activations, t_labels, t_input_lengths], outputs=outputs)
def test_sd_csc(): A = sp.sparse.rand(4, 5, density=0.60, format="csc", dtype=np.float32) b = np.random.rand(5, 2).astype(np.float32) target = A * b a_val = as_tensor_variable(A.data) a_ind = as_tensor_variable(A.indices) a_ptr = as_tensor_variable(A.indptr) nrows = as_tensor_variable(np.int32(A.shape[0])) b = as_tensor_variable(b) res = aesara.sparse.opt.sd_csc(a_val, a_ind, a_ptr, nrows, b).eval() utt.assert_allclose(res, target)
def make_node(self, x): x = aet.as_tensor_variable(x) self_axis = self.axis if self_axis is None: broadcastable = [False] else: if self_axis < 0: self_axis += len(x.broadcastable) if self_axis < 0 or self_axis >= len(x.broadcastable): raise RuntimeError( "Unique axis `{}` is outside of input ndim = " "{}.".format(self.axis, len(x.broadcastable))) broadcastable = [ b if axis != self_axis else False for axis, b in enumerate(x.broadcastable) ] outputs = [TensorType(broadcastable=broadcastable, dtype=x.dtype)()] typ = TensorType(broadcastable=[False], dtype="int64") if self.return_index: outputs.append(typ()) if self.return_inverse: outputs.append(typ()) if self.return_counts: outputs.append(typ()) return Apply(self, [x], outputs)
def shape_padaxis(t, axis): """Reshape `t` by inserting 1 at the dimension `axis`. Examples -------- >>> tensor = aesara.tensor.type.tensor3() >>> aesara.tensor.shape_padaxis(tensor, axis=0) DimShuffle{x,0,1,2}.0 >>> aesara.tensor.shape_padaxis(tensor, axis=1) DimShuffle{0,x,1,2}.0 >>> aesara.tensor.shape_padaxis(tensor, axis=3) DimShuffle{0,1,2,x}.0 >>> aesara.tensor.shape_padaxis(tensor, axis=-1) DimShuffle{0,1,2,x}.0 See Also -------- shape_padleft shape_padright Dimshuffle """ _t = aet.as_tensor_variable(t) ndim = _t.ndim + 1 if not -ndim <= axis < ndim: msg = "axis {0} is out of bounds [-{1}, {1})".format(axis, ndim) raise IndexError(msg) if axis < 0: axis += ndim pattern = [i for i in range(_t.type.ndim)] pattern.insert(axis, "x") return _t.dimshuffle(pattern)
def normalize_size_param(size): """Create an Aesara value for a ``RandomVariable`` ``size`` parameter.""" if size is None: size = constant([], dtype="int64") elif isinstance(size, int): size = as_tensor_variable([size], ndim=1) elif not isinstance(size, (np.ndarray, Variable, Sequence)): raise TypeError( "Parameter size must be None, an integer, or a sequence with integers." ) else: size = cast(as_tensor_variable(size, ndim=1), "int64") assert size.dtype in int_dtypes return size
def norm(x, ord): x = as_tensor_variable(x) ndim = x.ndim if ndim == 0: raise ValueError("'axis' entry is out of bounds.") elif ndim == 1: if ord is None: return tm.sum(x**2)**0.5 elif ord == "inf": return tm.max(abs(x)) elif ord == "-inf": return tm.min(abs(x)) elif ord == 0: return x[x.nonzero()].shape[0] else: try: z = tm.sum(abs(x**ord))**(1.0 / ord) except TypeError: raise ValueError("Invalid norm order for vectors.") return z elif ndim == 2: if ord is None or ord == "fro": return tm.sum(abs(x**2))**(0.5) elif ord == "inf": return tm.max(tm.sum(abs(x), 1)) elif ord == "-inf": return tm.min(tm.sum(abs(x), 1)) elif ord == 1: return tm.max(tm.sum(abs(x), 0)) elif ord == -1: return tm.min(tm.sum(abs(x), 0)) else: raise ValueError(0) elif ndim > 2: raise NotImplementedError("We don't support norm with ndim > 2")
def make_node(self, rng, size, dtype, *dist_params): """Create a random variable node. XXX: Unnamed/non-keyword arguments are considered distribution parameters! If you want to set `size`, `rng`, and/or `name`, use their keywords. Parameters ---------- rng: RandomStateType Existing Aesara `RandomState` object to be used. Creates a new one, if `None`. size: int or Sequence Numpy-like size of the output (i.e. replications). dtype: str The dtype of the sampled output. If the value ``"floatX"`` is given, then ``dtype`` is set to ``aesara.config.floatX``. This value is only used when `self.dtype` isn't set. dist_params: list Distribution parameters. Results ------- out: `Apply` A node with inputs `(rng, size, dtype) + dist_args` and outputs `(rng_var, out_var)`. """ size = normalize_size_param(size) dist_params = tuple( as_tensor_variable(p) if not isinstance(p, Variable) else p for p in dist_params ) if rng is None: rng = aesara.shared(np.random.RandomState()) elif not isinstance(rng.type, RandomStateType): raise TypeError("The type of rng should be an instance of RandomStateType") bcast = self.compute_bcast(dist_params, size) dtype = self.dtype or dtype if dtype == "floatX": dtype = config.floatX elif dtype is None or (isinstance(dtype, str) and dtype not in all_dtypes): raise TypeError("dtype is unspecified") if isinstance(dtype, str): dtype_idx = constant(all_dtypes.index(dtype), dtype="int64") else: dtype_idx = constant(dtype, dtype="int64") dtype = all_dtypes[dtype_idx.data] outtype = TensorType(dtype=dtype, broadcastable=bcast) out_var = outtype() inputs = (rng, size, dtype_idx) + dist_params outputs = (rng.type(), out_var) return Apply(self, inputs, outputs)
def make_node(self, condition, *monitored_vars): # Ensure that condition is an Aesara tensor if not isinstance(condition, Variable): condition = as_tensor_variable(condition) # Validate that the condition is a scalar (else it is not obvious how # is should be evaluated) assert condition.ndim == 0 # Because the user might be tempted to instantiate PdbBreakpoint only # once and apply it many times on different number of inputs, we must # create a new instance of the op here, define the instance attributes # (view_map and var_types) in that instance and then apply it on the # inputs. new_op = PdbBreakpoint(name=self.name) new_op.view_map = {} new_op.inp_types = [] for i in range(len(monitored_vars)): # Every output i is a view of the input i+1 because of the input # condition. new_op.view_map[i] = [i + 1] new_op.inp_types.append(monitored_vars[i].type) # Build the Apply node inputs = [condition] + list(monitored_vars) outputs = [inp.type() for inp in monitored_vars] return Apply(op=new_op, inputs=inputs, outputs=outputs)
def test_jax_scan_tap_output(): a_aet = scalar("a") def input_step_fn(y_tm1, y_tm3, a): y_tm1.name = "y_tm1" y_tm3.name = "y_tm3" res = (y_tm1 + y_tm3) * a res.name = "y_t" return res y_scan_aet, _ = scan( fn=input_step_fn, outputs_info=[ { "initial": aet.as_tensor_variable(np.r_[-1.0, 1.3, 0.0].astype(config.floatX)), "taps": [-1, -3], }, ], non_sequences=[a_aet], n_steps=10, name="y_scan", ) y_scan_aet.name = "y" y_scan_aet.owner.inputs[0].name = "y_all" out_fg = FunctionGraph([a_aet], [y_scan_aet]) test_input_vals = [np.array(10.0).astype(config.floatX)] compare_jax_and_py(out_fg, test_input_vals)
def make_node(self, inp, s=None): # A shape parameter s can be provided as an input. For now this is used to # manage odd transform sizes. # Later this could be extended to handle padding and trunkation, # following numpy's interface. However, cuFFT expects array that match # the shape given to the plan, so padding will have to be done in the op. # The effect of padding on gradients has yet to be investigated. if not skcuda_available: raise RuntimeError("skcuda is needed for CuFFTOp") if not pygpu_available: raise RuntimeError("pygpu is needed for CuFFTOp") if not pycuda_available: raise RuntimeError("pycuda is needed for CuFFTOp") inp = gpu_contiguous(as_gpuarray_variable(inp, infer_context_name(inp))) # If no shape is provided as input, default to input data shape. if s is None: s = inp.shape[1:] s = as_tensor_variable(s) assert inp.dtype == "float32" assert s.ndim == 1 assert s.dtype in integer_dtypes return Apply(self, [inp, s], [self.output_type(inp)()])
def make_node(self, a, s=None): a = as_tensor_variable(a) if a.ndim < 2: raise TypeError( "%s: input must have dimension > 2, with first dimension batches" % self.__class__.__name__) if s is None: s = a.shape[1:] s = as_tensor_variable(s) else: s = as_tensor_variable(s) if s.dtype not in integer_dtypes: raise TypeError("%s: length of the transformed axis must be" " of type integer" % self.__class__.__name__) return Apply(self, [a, s], [self.output_type(a)()])