Make scalar Categorical sampling work in recent versions of Numba

Author: brandonwillard

pytensor/link/numba/dispatch/random.py

@@ -5,7 +5,15 @@
 import numpy as np
 from numba import _helperlib, types
 from numba.core import cgutils
-from numba.extending import NativeValue, box, models, register_model, typeof_impl, unbox
+from numba.extending import (
+    NativeValue,
+    box,
+    models,
+    overload,
+    register_model,
+    typeof_impl,
+    unbox,
+)
 from numpy.random import RandomState
 
 import pytensor.tensor.random.basic as aer
@@ -78,6 +86,16 @@ def box_random_state(typ, val, c):
     return class_obj
 
 
+@overload(np.random.uniform)
+def uniform_empty_size(a, b, size):
+    if isinstance(size, types.Tuple) and size.count == 0:
+
+        def uniform_no_size(a, b, size):
+            return np.random.uniform(a, b)
+
+        return uniform_no_size
+
+
 @numba_typify.register(RandomState)
 def numba_typify_RandomState(state, **kwargs):
     # The numba_typify in this case is just an passthrough function
@@ -321,7 +339,7 @@ def categorical_rv(rng, size, dtype, p):
             size_tpl = numba_ndarray.to_fixed_tuple(size, size_len)
             p = np.broadcast_to(p, size_tpl + p.shape[-1:])
 
-        unif_samples = np.random.uniform(0, 1, size_tpl)
+        unif_samples = np.asarray(np.random.uniform(0, 1, size_tpl))
 
         res = np.empty(size_tpl, dtype=out_dtype)
         for idx in np.ndindex(*size_tpl):