Fix non-vm NUMBA

Author: ricardoV94

pytensor/link/numba/dispatch/basic.py

@@ -270,9 +270,15 @@ def numba_funcify_FunctionGraph(
     jit_nodes: bool = False,
     **kwargs,
 ):
+    def numba_funcify_wrapper(*args, **kwargs):
+        result = numba_funcify(*args, **kwargs)
+        if isinstance(result, tuple):
+            return result[0]
+        return result
+
     return fgraph_to_python(
         fgraph,
-        op_conversion_fn=numba_funcify_njit if jit_nodes else numba_funcify,
+        op_conversion_fn=numba_funcify_njit if jit_nodes else numba_funcify_wrapper,
         type_conversion_fn=numba_typify,
         fgraph_name=fgraph_name,
         **kwargs,

pytensor/link/numba/dispatch/subtensor.py

@@ -16,7 +16,6 @@
     Subtensor,
 )
 from pytensor.tensor.type_other import NoneTypeT, SliceType
-from pytensor.utils import hash_from_code
 
 
 @numba_funcify.register(Subtensor)
@@ -102,7 +101,7 @@ def {function_name}({", ".join(input_names)}):
         function_name=function_name,
         global_env=globals() | {"np": np},
     )
-    return numba_njit(func, boundscheck=True), hash_from_code(subtensor_def_src)
+    return numba_njit(func, boundscheck=True)
 
 
 @numba_funcify.register(AdvancedSubtensor)

pytensor/link/numba/dispatch/tensor_basic.py

@@ -110,10 +110,7 @@ def alloc(val, {", ".join(shape_var_names)}):
         {**globals(), **global_env},
     )
 
-    return (
-        pytensor.link.numba.compile.numba_njit(alloc_fn),
-        hash_from_code(alloc_def_src),
-    )
+    return pytensor.link.numba.compile.numba_njit(alloc_fn)
 
 
 @numba_funcify.register(ARange)

pytensor/link/numba/linker.py

@@ -19,7 +19,7 @@ def jit_compile(self, fn):
         else:
             from pytensor.link.numba.compile import numba_njit
 
-            jitted_fn = numba_njit(fn, no_cpython_wrapper=False, no_cfunc_wrapper=False)
+            jitted_fn = numba_njit(fn, final_function=True)
             return jitted_fn
 
     def create_thunk_inputs(self, storage_map):