Deal with single output scalar_ops

Author: aseyboldt

pytensor/link/numba/dispatch/elemwise.py

@@ -443,6 +443,7 @@ def _vectorize_bc(
     scalar_func,
     input_bc_patterns,
     output_bc_patterns,
+    output_dtypes,
     boundscheck=False,
     noalias_outputs=False,
 ):
@@ -484,8 +485,9 @@ def codegen(context, builder, signature, args):
             shape = cgutils.unpack_tuple(builder, iter_shape)
 
             # Lower the code of the scalar function so that we can use it in the inner loop
+            # Caching is set to false to avoid a numba bug TODO ref?
             inner = context.compile_subroutine(
-                builder, scalar_func, scalar_signature
+                builder, scalar_func, scalar_signature, caching=False,
             ).fndesc
 
             # Extract shape and stride information from the array.
@@ -546,9 +548,15 @@ def extract_array(aryty, ary):
 
             # Call scalar function
             output_values = context.call_internal(
-                builder, inner, scalar_signature, input_vals
+                builder,
+                inner,
+                scalar_signature,
+                input_vals,
             )
-            output_values = cgutils.unpack_tuple(builder, output_values)
+            if isinstance(scalar_signature.return_type, types.Tuple):
+                output_values = cgutils.unpack_tuple(builder, output_values)
+            else:
+                output_values = [output_values]
 
             # Update output value or accumulators respectively
             for i, ((accu, _), value) in enumerate(
@@ -614,9 +622,6 @@ def impl_vectorized(*inputs):
 
         iter_shape_repeated = tuple([iter_shape_template[:] for _ in range(n_outputs)])
 
-        # TODO Infer from signature
-        output_dtypes = (np.float64,) * n_outputs
-
         @numba.extending.register_jitable
         def make_output(iter_shape, bc, dtype):
             shape = iter_shape
@@ -684,19 +689,29 @@ def numba_funcify_Elemwise(op, node, **kwargs):
 
     assert not op.inplace_pattern
 
-    @register_jitable
-    def wrapper(in1, in2):
-        return (scalar_op_fn(in1, in2),)
+    #scalar_wrapper = register_jitable(scalar_op_fn)
+    scalar_wrapper = scalar_op_fn
 
     ndim = node.outputs[0].ndim
     output_bc_patterns = tuple([(False,) * ndim for _ in node.outputs])
     input_bc_patterns = tuple([input_var.broadcastable for input_var in node.inputs])
 
-    vectorized = _vectorize_bc(wrapper, input_bc_patterns, output_bc_patterns)
+    vectorized = _vectorize_bc(
+        scalar_wrapper,
+        input_bc_patterns,
+        output_bc_patterns,
+        output_dtypes=tuple([
+            variable.dtype
+            for variable in node.outputs
+        ]),
+    )
 
-    @numba_njit
-    def elemwise_wrapper(in1, in2):
-        return vectorized(in1, in2)[0]
+    if len(node.outputs) == 1:
+        @numba_njit
+        def elemwise_wrapper(*inputs):
+            return vectorized(*inputs)[0]
+    else:
+        elemwise_wrapper = vectorized
 
     return elemwise_wrapper