Cache the whole fgraph

Author: ricardoV94

pytensor/link/numba/dispatch/basic.py

@@ -275,7 +275,7 @@ def type_conversion_and_key_collection(value, variable, **kwargs):
         fgraph_key = sha256(
             str((tuple(cache_keys), len(fgraph.inputs), len(fgraph.outputs))).encode()
         ).hexdigest()
-    return py_func, fgraph_key
+    return numba_njit(py_func), fgraph_key
 
 
 @numba_funcify.register(OpFromGraph)

pytensor/link/numba/dispatch/scalar.py

@@ -238,7 +238,6 @@ def numba_funcify_Composite(op, node, **kwargs):
     composite_fn, fgraph_key = numba_funcify_and_cache_key(
         op.fgraph, squeeze_output=True, **kwargs
     )
-    composite_fn = numba_njit(composite_fn)
     if fgraph_key is None:
         composite_key = None
     else:

pytensor/link/numba/dispatch/scan.py

@@ -100,7 +100,7 @@ def numba_funcify_Scan(op: Scan, node, **kwargs):
     )
     rewriter(fgraph)
 
-    scan_inner_func = numba_njit(numba_funcify(op.fgraph))
+    scan_inner_func = numba_funcify(op.fgraph)
 
     outer_in_names_to_vars = {
         (f"outer_in_{i}" if i > 0 else "n_steps"): v for i, v in enumerate(node.inputs)

pytensor/link/numba/linker.py

@@ -10,18 +10,12 @@ def __init__(self, *args, vm: bool = False, **kwargs):
 
     def fgraph_convert(self, fgraph, **kwargs):
         import pytensor.link.numba.dispatch  # noqa
-        from pytensor.link.numba.compile import numba_funcify
+        from pytensor.link.numba.cache import numba_njit_and_cache
 
-        return numba_funcify(fgraph, **kwargs)
+        return numba_njit_and_cache(fgraph, **kwargs)[0]
 
     def jit_compile(self, fn):
-        if self.vm:
-            return fn
-        else:
-            from pytensor.link.numba.cache import numba_njit
-
-            jitted_fn = numba_njit(fn, final_function=True)
-            return jitted_fn
+        return fn
 
     def create_thunk_inputs(self, storage_map):
         return [storage_map[n] for n in self.fgraph.inputs]