1import random 2 3import numpy as np 4 5from numba.core import types 6from .templates import (ConcreteTemplate, AbstractTemplate, AttributeTemplate, 7 CallableTemplate, Registry, signature) 8from numba.np.numpy_support import numpy_version 9 10 11registry = Registry() 12infer = registry.register 13infer_global = registry.register_global 14infer_getattr = registry.register_attr 15 16 17# random.random(), random.seed() etc. are not plain functions, they are bound 18# methods of a private object. We have to be careful to use a well-known 19# object (e.g. the string "random.seed") as a key, not the bound method itself. 20# (same for np.random.random(), etc.) 21 22_int_types = sorted(set((types.intp, types.int64))) 23# Should we support float32? 24_float_types = [types.float64] 25 26 27# 28# Basics 29# 30 31def normalize_shape(shape): 32 if isinstance(shape, types.Integer): 33 return types.intp, 1 34 elif (isinstance(shape, types.BaseTuple) and 35 all(isinstance(v, types.Integer)) for v in shape): 36 ndim = len(shape) 37 return types.UniTuple(types.intp, ndim), ndim 38 else: 39 raise TypeError("invalid size type %s" % (shape,)) 40 41 42class RandomTemplate(CallableTemplate): 43 """ 44 A template helper to transparently handle the typing of array-returning 45 np.random.* functions. 46 """ 47 48 def array_typer(self, scalar_typer, size=None): 49 prefix = self.key.split('.')[0] 50 assert prefix in ('np', 'random'), self.key 51 52 if size is None: 53 # Scalar variant 54 def typer(*args, **kwargs): 55 return scalar_typer(*args, **kwargs) 56 else: 57 # Array variant (only for the 'np.random.*' namespace) 58 def typer(*args, **kwargs): 59 if prefix == 'random': 60 raise TypeError("unexpected size parameter for %r" 61 % (self.key,)) 62 shape, ndim = normalize_shape(size) 63 # Type the scalar variant and wrap the result in an array 64 # of the appropriate dimensionality. 65 sig = scalar_typer(*args, **kwargs) 66 if sig is not None: 67 return signature( 68 types.Array(sig.return_type, ndim, 'C'), 69 *(sig.args + (shape,))) 70 71 return typer 72 73 74class ConcreteRandomTemplate(RandomTemplate): 75 """ 76 A RandomTemplate subclass using the `cases` attribute as a list of 77 allowed scalar signatures. 78 """ 79 80 def array_typer(self, size=None): 81 key = self.key 82 cases = self.cases 83 context = self.context 84 85 def concrete_scalar_typer(*args, **kwargs): 86 # Filter out omitted args 87 while args and args[-1] is None: 88 args = args[:-1] 89 return context.resolve_overload(key, cases, args, kwargs) 90 91 return RandomTemplate.array_typer(self, concrete_scalar_typer, size) 92 93 94@infer_global(random.getrandbits, typing_key="random.getrandbits") 95class Random_getrandbits(ConcreteTemplate): 96 cases = [signature(types.uint64, types.int32)] 97 98@infer_global(random.random, typing_key="random.random") 99@infer_global(np.random.random, typing_key="np.random.random") 100class Random_random(ConcreteRandomTemplate): 101 cases = [signature(types.float64)] 102 103 def generic(self): 104 def typer(size=None): 105 return self.array_typer(size)() 106 return typer 107 108if numpy_version >= (1, 17): 109 infer_global( 110 np.random.random_sample, 111 typing_key="np.random.random_sample", 112 )(Random_random) 113 infer_global( 114 np.random.sample, 115 typing_key="np.random.sample", 116 )(Random_random) 117 infer_global( 118 np.random.ranf, 119 typing_key="np.random.ranf", 120 )(Random_random) 121 122 123@infer_global(random.randint, typing_key="random.randint") 124class Random_randint(ConcreteTemplate): 125 cases = [signature(tp, tp, tp) for tp in _int_types] 126 127@infer_global(np.random.randint, typing_key="np.random.randint") 128class Random_randint(ConcreteRandomTemplate): 129 cases = [signature(tp, tp) for tp in _int_types] 130 cases += [signature(tp, tp, tp) for tp in _int_types] 131 132 def generic(self): 133 def typer(low, high=None, size=None): 134 return self.array_typer(size)(low, high) 135 return typer 136 137 138@infer_global(random.randrange, typing_key="random.randrange") 139class Random_randrange(ConcreteTemplate): 140 cases = [signature(tp, tp) for tp in _int_types] 141 cases += [signature(tp, tp, tp) for tp in _int_types] 142 cases += [signature(tp, tp, tp, tp) for tp in _int_types] 143 144@infer_global(random.seed, typing_key="random.seed") 145@infer_global(np.random.seed, typing_key="np.random.seed") 146class Random_seed(ConcreteTemplate): 147 cases = [signature(types.void, types.uint32)] 148 149 150# 151# Distributions 152# 153 154@infer_global(np.random.geometric, typing_key="np.random.geometric") 155@infer_global(np.random.logseries, typing_key="np.random.logseries") 156@infer_global(np.random.zipf, typing_key="np.random.zipf") 157class Numpy_geometric(ConcreteRandomTemplate): 158 cases = [signature(types.int64, tp) for tp in _float_types] 159 160 def generic(self): 161 def typer(a, size=None): 162 return self.array_typer(size)(a) 163 return typer 164 165@infer_global(np.random.binomial, typing_key="np.random.binomial") 166@infer_global(np.random.negative_binomial, 167 typing_key="np.random.negative_binomial") 168class Numpy_negative_binomial(ConcreteRandomTemplate): 169 cases = [signature(types.int64, types.int64, tp) for tp in _float_types] 170 171 def generic(self): 172 def typer(n, p, size=None): 173 return self.array_typer(size)(n, p) 174 return typer 175 176@infer_global(np.random.poisson, typing_key="np.random.poisson") 177class Numpy_poisson(ConcreteRandomTemplate): 178 cases = [signature(types.int64, tp) for tp in _float_types] 179 cases += [signature(types.int64)] 180 181 def generic(self): 182 def typer(lam=None, size=None): 183 return self.array_typer(size)(lam) 184 return typer 185 186@infer_global(np.random.exponential, typing_key="np.random.exponential") 187@infer_global(np.random.rayleigh, typing_key="np.random.rayleigh") 188class Numpy_exponential(ConcreteRandomTemplate): 189 cases = [signature(tp, tp) for tp in _float_types] 190 cases += [signature(tp) for tp in _float_types] 191 192 def generic(self): 193 def typer(scale=None, size=None): 194 return self.array_typer(size)(scale) 195 return typer 196 197@infer_global(np.random.hypergeometric, typing_key="np.random.hypergeometric") 198class Numpy_hypergeometric(ConcreteRandomTemplate): 199 cases = [signature(tp, tp, tp, tp) for tp in _int_types] 200 201 def generic(self): 202 def typer(ngood, nbad, nsample, size=None): 203 return self.array_typer(size)(ngood, nbad, nsample) 204 return typer 205 206@infer_global(np.random.laplace, typing_key="np.random.laplace") 207@infer_global(np.random.logistic, typing_key="np.random.logistic") 208@infer_global(np.random.lognormal, typing_key="np.random.lognormal") 209@infer_global(np.random.normal, typing_key="np.random.normal") 210class Numpy_normal(ConcreteRandomTemplate): 211 cases = [signature(tp, tp, tp) for tp in _float_types] 212 cases += [signature(tp, tp) for tp in _float_types] 213 cases += [signature(tp) for tp in _float_types] 214 215 def generic(self): 216 def typer(loc=None, scale=None, size=None): 217 return self.array_typer(size)(loc, scale) 218 return typer 219 220@infer_global(np.random.gamma, typing_key="np.random.gamma") 221class Numpy_gamma(ConcreteRandomTemplate): 222 cases = [signature(tp, tp, tp) for tp in _float_types] 223 cases += [signature(tp, tp) for tp in _float_types] 224 225 def generic(self): 226 def typer(shape, scale=None, size=None): 227 return self.array_typer(size)(shape, scale) 228 return typer 229 230@infer_global(np.random.triangular, typing_key="np.random.triangular") 231class Random_ternary_distribution(ConcreteRandomTemplate): 232 cases = [signature(tp, tp, tp, tp) for tp in _float_types] 233 234 def generic(self): 235 def typer(left, mode, right, size=None): 236 return self.array_typer(size)(left, mode, right) 237 return typer 238 239 240@infer_global(np.random.beta, typing_key="np.random.beta") 241@infer_global(np.random.f, typing_key="np.random.f") 242@infer_global(np.random.gumbel, typing_key="np.random.gumbel") 243@infer_global(np.random.uniform, typing_key="np.random.uniform") 244@infer_global(np.random.vonmises, typing_key="np.random.vonmises") 245@infer_global(np.random.wald, typing_key="np.random.wald") 246@infer_global(random.betavariate, typing_key="random.betavariate") 247@infer_global(random.gammavariate, typing_key="random.gammavariate") 248@infer_global(random.gauss, typing_key="random.gauss") 249@infer_global(random.lognormvariate, typing_key="random.lognormvariate") 250@infer_global(random.normalvariate, typing_key="random.normalvariate") 251@infer_global(random.uniform, typing_key="random.uniform") 252@infer_global(random.vonmisesvariate, typing_key="random.vonmisesvariate") 253@infer_global(random.weibullvariate, typing_key="random.weibullvariate") 254class Random_binary_distribution(ConcreteRandomTemplate): 255 cases = [signature(tp, tp, tp) for tp in _float_types] 256 257 def generic(self): 258 def typer(a, b, size=None): 259 return self.array_typer(size)(a, b) 260 return typer 261 262 263@infer_global(np.random.chisquare, typing_key="np.random.chisquare") 264@infer_global(np.random.pareto, typing_key="np.random.pareto") 265@infer_global(np.random.power, typing_key="np.random.power") 266@infer_global(np.random.standard_gamma, typing_key="np.random.standard_gamma") 267@infer_global(np.random.standard_t, typing_key="np.random.standard_t") 268@infer_global(np.random.weibull, typing_key="np.random.weibull") 269@infer_global(random.expovariate, typing_key="random.expovariate") 270@infer_global(random.paretovariate, typing_key="random.paretovariate") 271class Random_unary_distribution(ConcreteRandomTemplate): 272 cases = [signature(tp, tp) for tp in _float_types] 273 274 def generic(self): 275 def typer(a, size=None): 276 return self.array_typer(size)(a) 277 return typer 278 279 280@infer_global(np.random.standard_cauchy, 281 typing_key="np.random.standard_cauchy") 282@infer_global(np.random.standard_normal, 283 typing_key="np.random.standard_normal") 284@infer_global(np.random.standard_exponential, 285 typing_key="np.random.standard_exponential") 286class Random_nullary_distribution(ConcreteRandomTemplate): 287 cases = [signature(tp) for tp in _float_types] 288 289 def generic(self): 290 def typer(size=None): 291 return self.array_typer(size)() 292 return typer 293 294 295@infer_global(random.triangular, typing_key="random.triangular") 296class Random_triangular(ConcreteTemplate): 297 cases = [signature(tp, tp, tp) for tp in _float_types] 298 cases += [signature(tp, tp, tp, tp) for tp in _float_types] 299 300# NOTE: some functions can have @overloads in numba.targets.randomimpl, 301# and therefore don't need a typing declaration here. 302