Break things to understand vector inputs

Author: ricardoV94

pytensor/link/numba/dispatch/random.py

@@ -181,10 +181,10 @@ def numba_funcify_RandomVariable(op: RandomVariable, node, **kwargs):
     if op.ndim_supp > 0:
         raise NotImplementedError("Multivariate random variables not supported yet")
 
-    if any(ndim_param > 0 for ndim_param in op.ndims_params):
-        raise NotImplementedError(
-            "Random variables with non scalar core inputs not supported yet"
-        )
+    # if any(ndim_param > 0 for ndim_param in op.ndims_params):
+    #     raise NotImplementedError(
+    #         "Random variables with non scalar core inputs not supported yet"
+    #     )
 
     # TODO: Use dispatch, so users can define the core case
     #  Use string repr for default like below
@@ -197,14 +197,19 @@ def numba_funcify_RandomVariable(op: RandomVariable, node, **kwargs):
     # exec(inner_code)
     # scalar_op_fn = locals()['scalar_op_fn']
 
+    # @numba_basic.numba_njit
+    # def scalar_op_fn(rng, mu, scale):
+    #     return rng.normal(mu, scale)
+
     @numba_basic.numba_njit
-    def scalar_op_fn(rng, mu, scale):
-        return rng.normal(mu, scale)
+    def scalar_op_fn(rng, p):
+        unif_sample = rng.uniform(0, 1)
+        return np.searchsorted(np.cumsum(p), unif_sample)
 
-    ndim = node.default_output().ndim
-    output_bc_patterns = ((False,) * ndim,)
+    batch_ndim = node.default_output().ndim - op.ndim_supp
+    output_bc_patterns = ((False,) * batch_ndim,)
     input_bc_patterns = tuple(
-        [input_var.broadcastable for input_var in node.inputs[3:]]
+        [input_var.broadcastable[:batch_ndim] for input_var in node.inputs[3:]]
     )
     output_dtypes = (node.default_output().type.dtype,)
     inplace_pattern = ()  # tuple(op.inplace_pattern.items())
@@ -325,7 +330,7 @@ def body_fn(a):
     )
 
 
-@numba_funcify.register(ptr.CategoricalRV)
+# @numba_funcify.register(ptr.CategoricalRV)
 def numba_funcify_CategoricalRV(op, node, **kwargs):
     out_dtype = node.outputs[1].type.numpy_dtype
     size_len = int(get_vector_length(node.inputs[1]))

pytensor/link/numba/dispatch/vectorize_codegen.py

@@ -12,6 +12,7 @@
 from numba.core.base import BaseContext
 from numba.core.types.misc import NoneType
 from numba.np import arrayobj
+from numba.np.ufunc.wrappers import _ArrayArgLoader
 
 
 def compute_itershape(
@@ -22,11 +23,11 @@ def compute_itershape(
     size: list[ir.Instruction] | None,
 ):
     one = ir.IntType(64)(1)
-    ndim = len(in_shapes[0])
-    shape = [None] * ndim
+    batch_ndim = len(broadcast_pattern[0])
+    shape = [None] * batch_ndim
     if size is not None:
         shape = size
-        for i in range(ndim):
+        for i in range(batch_ndim):
             for j, (bc, in_shape) in enumerate(zip(broadcast_pattern, in_shapes)):
                 length = in_shape[i]
                 if bc[i]:
@@ -61,7 +62,7 @@ def compute_itershape(
                                 )
     else:
         # Size is implied by the broadcast pattern
-        for i in range(ndim):
+        for i in range(batch_ndim):
             for j, (bc, in_shape) in enumerate(zip(broadcast_pattern, in_shapes)):
                 length = in_shape[i]
                 if bc[i]:
@@ -96,7 +97,7 @@ def compute_itershape(
                                 )
                 else:
                     shape[i] = length
-        for i in range(ndim):
+        for i in range(batch_ndim):
             if shape[i] is None:
                 shape[i] = one
     return shape
@@ -157,7 +158,7 @@ def make_loop_call(
     input_types: tuple[Any, ...],
     output_types: tuple[Any, ...],
 ):
-    safe = (False, False)
+    # safe = (False, False)
 
     n_outputs = len(outputs)
 
@@ -182,6 +183,12 @@ def extract_array(aryty, obj):
     # input_scope_set = mod.add_metadata([input_scope, output_scope])
     # output_scope_set = mod.add_metadata([input_scope, output_scope])
 
+    typ = input_types[0]
+    inp = inputs[0]
+    shape = cgutils.unpack_tuple(builder, inp.shape)
+    strides = cgutils.unpack_tuple(builder, inp.strides)
+    loader = _ArrayArgLoader(typ.dtype, typ.ndim, shape[-1], False, shape, strides)
+
     inputs = tuple(extract_array(aryty, ary) for aryty, ary in zip(input_types, inputs))
 
     outputs = tuple(
@@ -216,8 +223,9 @@ def extract_array(aryty, obj):
     input_vals = []
     for array_info, bc in zip(inputs, input_bc):
         idxs_bc = [zero if bc else idx for idx, bc in zip(idxs, bc)]
-        ptr = cgutils.get_item_pointer2(context, builder, *array_info, idxs_bc, *safe)
-        val = builder.load(ptr)
+        # ptr = cgutils.get_item_pointer2(context, builder, *array_info, idxs_bc, *safe)
+        val = loader.load(context, builder, inp.data, idxs[0] or zero)
+        # val = builder.load(ptr)
         # val.set_metadata("alias.scope", input_scope_set)
         # val.set_metadata("noalias", output_scope_set)
         input_vals.append(val)
@@ -340,16 +348,21 @@ def _vectorized(
     if not all(isinstance(input, types.Array) for input in inputs):
         raise TypingError("Vectorized inputs must be arrays.")
 
-    ndim = inputs[0].ndim
+    batch_ndim = len(input_bc_patterns[0])
 
-    if not all(input.ndim == ndim for input in inputs):
+    if not all(input.ndim >= batch_ndim for input in inputs):
         raise TypingError("Vectorized inputs must have the same rank.")
 
-    if not all(len(pattern) == ndim for pattern in output_bc_patterns):
+    if not all(len(pattern) >= batch_ndim for pattern in output_bc_patterns):
         raise TypingError("Invalid output broadcasting pattern.")
 
     scalar_signature = typingctx.resolve_function_type(
-        scalar_func, [*constant_inputs, *[in_type.dtype for in_type in inputs]], {}
+        scalar_func,
+        [
+            *constant_inputs,
+            *[in_type.dtype if in_type.ndim == 0 else in_type for in_type in inputs],
+        ],
+        {},
     )
 
     # So we can access the constant values in codegen...
@@ -430,7 +443,7 @@ def codegen(
         )
 
     ret_types = [
-        types.Array(numba.from_dtype(np.dtype(dtype)), ndim, "C")
+        types.Array(numba.from_dtype(np.dtype(dtype)), batch_ndim, "C")
         for dtype in output_dtypes
     ]
 

tests/link/numba/test_random.py

@@ -642,3 +642,10 @@ def test_rng_non_default_update():
     ref = np.random.default_rng(2).normal(size=10)
     np.testing.assert_allclose(fn(), ref)
     np.testing.assert_allclose(fn(), ref)
+
+
+def test_categorical_rv():
+    x = pt.random.categorical(p=[[0.5, 0, 0, 0.5], [0, 0.5, 0.5, 0]], size=(2,))
+    updates = {x.owner.inputs[0]: x.owner.outputs[0]}
+    fn = function([], x, updates=updates, mode="NUMBA")
+    print([fn() for _ in range(50)])