Saner defaults

Author: ricardoV94

pytensor/link/numba/cache.py

@@ -1,3 +1,4 @@
+import warnings
 import weakref
 from collections.abc import Callable
 from functools import singledispatch, wraps
@@ -25,8 +26,6 @@ def __init__(self, py_func, py_file, hash):
         self._py_func = py_func
         self._py_file = py_file
         self._hash = hash
-        # src_hash = hash(pytensor_loader._module_sources[self._py_file])
-        # self._hash = hash((src_hash, py_file, pytensor.__version__))
 
     def ensure_cache_path(self):
         pass
@@ -185,27 +184,34 @@ def dispatch_func_wrapper(*args, **kwargs):
     return decorator
 
 
-def numba_njit_and_cache(op, node, **kwargs):
-    jitable_func, key = numba_funcify_and_cache_key(op, node=node, **kwargs)
+def numba_njit_and_cache(op, *args, **kwargs):
+    jitable_func, key = numba_funcify_and_cache_key(op, *args, **kwargs)
 
     if key is not None:
         # To force numba to use our cache, we must compile the function so that any closure
         # becomes a global variable...
         op_name = op.__class__.__name__
-        cached_func = compile_and_cache_numba_function_src(
+        cached_func = compile_numba_function_src(
             src=f"def {op_name}(*args): return jitable_func(*args)",
             function_name=op_name,
-            global_env=globals() | dict(jitable_func=jitable_func),
+            global_env=globals() | {"jitable_func": jitable_func},
             cache_key=key,
         )
-        return numba_njit(cached_func, final_function=True, cache=True)
+        return numba_njit(cached_func, final_function=True, cache=True), key
     else:
+        if config.numba__cache and config.compiler_verbose:
+            warnings.warn(
+                f"Custom numba cache disabled for {op} of type {type(op)}. "
+                f"Even if the function is cached by numba, larger graphs using this function cannot be cached.\n"
+                "To enable custom caching, register a numba_funcify_and_cache_key implementation for this Op, with a proper cache key."
+            )
+
         return numba_njit(
             lambda *args: jitable_func(*args), final_function=True, cache=False
-        )
+        ), None
 
 
-def compile_and_cache_numba_function_src(
+def compile_numba_function_src(
     src: str,
     function_name: str,
     global_env: dict[Any, Any] | None = None,

pytensor/link/numba/dispatch/basic.py

@@ -3,6 +3,8 @@
 import warnings
 from collections.abc import Callable
 from functools import singledispatch
+from hashlib import sha256
+from pickle import dumps
 
 import numba
 import numpy as np
@@ -17,11 +19,13 @@
 from pytensor.compile.builders import OpFromGraph
 from pytensor.compile.function.types import add_supervisor_to_fgraph
 from pytensor.compile.ops import DeepCopyOp
+from pytensor.graph import Constant
 from pytensor.graph.fg import FunctionGraph
 from pytensor.graph.op import Op
 from pytensor.ifelse import IfElse
 from pytensor.link.numba.cache import (
     numba_njit_and_cache,
+    register_funcify_and_cache_key,
     register_funcify_default_op_cache_key,
 )
 from pytensor.link.numba.compile import (
@@ -226,21 +230,52 @@ def numba_funcify_fallback(
     return generate_fallback_impl(op, node, storage_map, **kwargs)
 
 
-@numba_funcify.register(FunctionGraph)
+def key_for_constant(data):
+    """Create a cache key for a constant value."""
+    # TODO: This is just a placeholder
+    if isinstance(data, (int | float | bool | type(None))):
+        return str(data)
+    try:
+        # For NumPy arrays
+        return sha256(data.tobytes()).hexdigest()
+    except AttributeError:
+        # Fallback for other types
+        return sha256(dumps(data)).hexdigest()
+
+
+@register_funcify_and_cache_key(FunctionGraph)
 def numba_funcify_FunctionGraph(
     fgraph,
     node=None,
     fgraph_name="numba_funcified_fgraph",
     **kwargs,
 ):
-    # TODO: Create hash key for whole graph
-    return fgraph_to_python(
+    cache_keys = []
+
+    def op_conversion_and_key_collection(*args, **kwargs):
+        func, key = numba_njit_and_cache(*args, **kwargs)
+        cache_keys.append(key)
+        return func
+
+    def type_conversion_and_key_collection(value, variable, **kwargs):
+        if isinstance(variable, Constant):
+            cache_keys.append(key_for_constant(value))
+        return numba_typify(value, variable=variable, **kwargs)
+
+    py_func = fgraph_to_python(
         fgraph,
-        op_conversion_fn=numba_njit_and_cache,
-        type_conversion_fn=numba_typify,
+        op_conversion_fn=op_conversion_and_key_collection,
+        type_conversion_fn=type_conversion_and_key_collection,
         fgraph_name=fgraph_name,
         **kwargs,
     )
+    if any(key is None for key in cache_keys):
+        fgraph_key = None
+    else:
+        fgraph_key = sha256(
+            str((tuple(cache_keys), len(fgraph.inputs), len(fgraph.outputs))).encode()
+        ).hexdigest()
+    return py_func, fgraph_key
 
 
 @numba_funcify.register(OpFromGraph)

pytensor/link/numba/dispatch/elemwise.py

@@ -9,7 +9,7 @@
 
 from pytensor.graph.op import Op
 from pytensor.link.numba.cache import (
-    compile_and_cache_numba_function_src,
+    compile_numba_function_src,
     numba_funcify_and_cache_key,
     register_funcify_and_cache_key,
     register_funcify_default_op_cache_key,
@@ -239,7 +239,7 @@ def {careduce_fn_name}(x):
     careduce_def_src += "\n\n"
     careduce_def_src += indent(f"return {return_obj}", " " * 4)
 
-    careduce_fn = compile_and_cache_numba_function_src(
+    careduce_fn = compile_numba_function_src(
         careduce_def_src, careduce_fn_name, {**globals(), **global_env}
     )
 

pytensor/link/numba/dispatch/scalar.py

@@ -6,7 +6,8 @@
 from pytensor.compile.ops import TypeCastingOp
 from pytensor.graph.basic import Variable
 from pytensor.link.numba.cache import (
-    compile_and_cache_numba_function_src,
+    compile_numba_function_src,
+    numba_funcify_and_cache_key,
     register_funcify_and_cache_key,
 )
 from pytensor.link.numba.compile import (
@@ -138,7 +139,7 @@ def {scalar_op_fn_name}({', '.join(input_names)}):
     return direct_cast(scalar_func_numba({converted_call_args}, np.intc(1)), output_dtype)
             """
 
-    scalar_op_fn = compile_and_cache_numba_function_src(
+    scalar_op_fn = compile_numba_function_src(
         scalar_op_src,
         scalar_op_fn_name,
         {**globals(), **global_env},
@@ -172,7 +173,7 @@ def binary_to_nary_func(inputs: list[Variable], binary_op_name: str, binary_op:
 def {binary_op_name}({input_signature}):
     return {output_expr}
     """
-    nary_fn = compile_and_cache_numba_function_src(nary_src, binary_op_name, globals())
+    nary_fn = compile_numba_function_src(nary_src, binary_op_name, globals())
 
     return nary_fn
 
@@ -234,8 +235,15 @@ def clip(x, min_val, max_val):
 def numba_funcify_Composite(op, node, **kwargs):
     _ = kwargs.pop("storage_map", None)
 
-    composite_fn = numba_njit(numba_funcify(op.fgraph, squeeze_output=True, **kwargs))
-    return composite_fn, str(tuple(type(node.op) for node in op.fgraph.toposort()))
+    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:
+        composite_key = sha256(str((type(op), fgraph_key)).encode()).hexdigest()
+    return composite_fn, composite_key
 
 
 @register_funcify_and_cache_key(Second)

pytensor/link/numba/dispatch/scan.py

@@ -7,7 +7,7 @@
 from pytensor import In
 from pytensor.compile.function.types import add_supervisor_to_fgraph
 from pytensor.compile.mode import NUMBA, get_mode
-from pytensor.link.numba.cache import compile_and_cache_numba_function_src
+from pytensor.link.numba.cache import compile_numba_function_src
 from pytensor.link.numba.compile import (
     create_arg_string,
     create_tuple_string,
@@ -443,12 +443,10 @@ def scan({", ".join(outer_in_names)}):
     }
     global_env["np"] = np
 
-    scan_op_fn = compile_and_cache_numba_function_src(
+    scan_op_fn = compile_numba_function_src(
         scan_op_src,
         "scan",
         {**globals(), **global_env},
-        # We can't cache until we can hash FunctionGraph
-        cache_key=None,
     )
 
-    return numba_njit(scan_op_fn, boundscheck=False), None
+    return numba_njit(scan_op_fn, boundscheck=False)

pytensor/link/numba/dispatch/subtensor.py

@@ -2,7 +2,7 @@
 
 from pytensor.graph import Type
 from pytensor.link.numba.cache import (
-    compile_and_cache_numba_function_src,
+    compile_numba_function_src,
     register_funcify_default_op_cache_key,
 )
 from pytensor.link.numba.compile import numba_njit
@@ -99,7 +99,7 @@ def {function_name}({", ".join(input_names)}):
     return np.asarray(z)
     """
 
-    func = compile_and_cache_numba_function_src(
+    func = compile_numba_function_src(
         subtensor_def_src,
         function_name=function_name,
         global_env=globals() | {"np": np},

pytensor/link/numba/dispatch/tensor_basic.py

@@ -3,7 +3,7 @@
 import numpy as np
 
 from pytensor.link.numba.cache import (
-    compile_and_cache_numba_function_src,
+    compile_numba_function_src,
     register_funcify_default_op_cache_key,
 )
 from pytensor.link.numba.compile import (
@@ -56,7 +56,7 @@ def allocempty({", ".join(shape_var_names)}):
     return np.empty(scalar_shape, dtype)
     """
 
-    alloc_fn = compile_and_cache_numba_function_src(
+    alloc_fn = compile_numba_function_src(
         alloc_def_src, "allocempty", {**globals(), **global_env}
     )
 
@@ -100,7 +100,7 @@ def alloc(val, {", ".join(shape_var_names)}):
     res[...] = val
     return res
     """
-    alloc_fn = compile_and_cache_numba_function_src(
+    alloc_fn = compile_numba_function_src(
         alloc_def_src,
         "alloc",
         {**globals(), **global_env},
@@ -223,7 +223,7 @@ def makevector({", ".join(input_names)}):
     return np.array({create_list_string(input_names)}, dtype=dtype)
     """
 
-    makevector_fn = compile_and_cache_numba_function_src(
+    makevector_fn = compile_numba_function_src(
         makevector_def_src,
         "makevector",
         {**globals(), **global_env},

pytensor/link/numba/dispatch/vectorize_codegen.py

@@ -15,7 +15,7 @@
 from numba.core.types.misc import NoneType
 from numba.np import arrayobj
 
-from pytensor.link.numba.cache import compile_and_cache_numba_function_src
+from pytensor.link.numba.cache import compile_numba_function_src
 from pytensor.link.numba.compile import numba_njit
 
 
@@ -53,7 +53,7 @@ def store_core_outputs({inp_signature}, {out_signature}):
 """
     global_env = {"core_op_fn": core_op_fn}
 
-    func = compile_and_cache_numba_function_src(
+    func = compile_numba_function_src(
         func_src,
         "store_core_outputs",
         {**globals(), **global_env},

pytensor/link/numba/linker.py

@@ -9,15 +9,15 @@ def __init__(self, *args, vm: bool = False, **kwargs):
         self.vm = vm
 
     def fgraph_convert(self, fgraph, **kwargs):
-        from pytensor.link.numba.dispatch import numba_funcify
+        from pytensor.link.numba.compile import numba_funcify
 
         return numba_funcify(fgraph, **kwargs)
 
     def jit_compile(self, fn):
         if self.vm:
             return fn
         else:
-            from pytensor.link.numba.compile import numba_njit
+            from pytensor.link.numba.cache import numba_njit
 
             jitted_fn = numba_njit(fn, final_function=True)
             return jitted_fn

tests/compile/function/test_types.py

@@ -1466,7 +1466,7 @@ def zerosumnormal(name, *, sigma=1.0, size, model_logp):
     return joined_inputs, [model_logp, model_dlogp]
 
 
-@pytest.mark.parametrize("mode", ["NUMBA", "C", "C_VM", "NUMBA"][:1])
+@pytest.mark.parametrize("mode", ["C", "C_VM", "NUMBA"][2:])
 def test_radon_model_compile_benchmark(mode, radon_model, benchmark):
     joined_inputs, [model_logp, model_dlogp] = radon_model
     rng = np.random.default_rng(1)