Fix numba FunctionGraph cache key

It's necessary to encode the edge information, not only the nodes and their ordering

Author: ricardoV94

pytensor/link/numba/dispatch/basic.py

@@ -10,7 +10,7 @@
 
 from pytensor import config
 from pytensor.graph.basic import Apply, Constant
-from pytensor.graph.fg import FunctionGraph
+from pytensor.graph.fg import FunctionGraph, Output
 from pytensor.graph.type import Type
 from pytensor.link.numba.cache import compile_numba_function_src, hash_from_pickle_dump
 from pytensor.link.numba.dispatch.sparse import CSCMatrixType, CSRMatrixType
@@ -501,28 +501,44 @@ def numba_funcify_FunctionGraph(
     cache_keys = []
     toposort = fgraph.toposort()
     clients = fgraph.clients
-    toposort_indices = {node: i for i, node in enumerate(toposort)}
-    # Add dummy output clients which are not included of the toposort
+    toposort_indices: dict[Apply | None, int] = {
+        node: i for i, node in enumerate(toposort)
+    }
+    # Use -1 for root inputs / constants whose owner is None
+    toposort_indices[None] = -1
+    # Add dummy output nodes which are not included of the toposort
     toposort_indices |= {
-        clients[out][0][0]: i
-        for i, out in enumerate(fgraph.outputs, start=len(toposort))
+        out_node: i + len(toposort)
+        for i, out in enumerate(fgraph.outputs)
+        for out_node, _ in clients[out]
+        if isinstance(out_node.op, Output) and out_node.op.idx == i
     }
 
-    def op_conversion_and_key_collection(*args, **kwargs):
+    def op_conversion_and_key_collection(op, *args, node, **kwargs):
         # Convert an Op to a funcified function and store the cache_key
 
         # We also Cache each Op so Numba can do less work next time it sees it
-        func, key = numba_funcify_ensure_cache(*args, **kwargs)
-        cache_keys.append(key)
+        func, key = numba_funcify_ensure_cache(op, node=node, *args, **kwargs)
+        if key is None:
+            cache_keys.append(key)
+        else:
+            # Add graph coordinate information (input edges and node location)
+            cache_keys.append(
+                (
+                    toposort_indices[node],
+                    tuple(toposort_indices[inp.owner] for inp in node.inputs),
+                    key,
+                )
+            )
         return func
 
     def type_conversion_and_key_collection(value, variable, **kwargs):
         # Convert a constant type to a numba compatible one and compute a cache key for it
 
-        # We need to know where in the graph the constants are used
-        # Otherwise we would hash stack(x, 5.0, 7.0), and stack(5.0, x, 7.0) the same
+        # Add graph coordinate information (client edges)
         # FIXME: It doesn't make sense to call type_conversion on non-constants,
-        #  but that's what fgraph_to_python currently does. We appease it, but don't consider for caching
+        #  but that's what fgraph_to_python currently does.
+        #  We appease it, but don't consider for caching
         if isinstance(variable, Constant):
             client_indices = tuple(
                 (toposort_indices[node], inp_idx) for node, inp_idx in clients[variable]
@@ -541,8 +557,24 @@ def type_conversion_and_key_collection(value, variable, **kwargs):
         # If a single element couldn't be cached, we can't cache the whole FunctionGraph either
         fgraph_key = None
     else:
+        # Add graph coordinate information for fgraph inputs (client edges) and fgraph outputs (input edges)
+        # Constant edges are handled by `type_conversion_and_key_collection` called by `fgraph_to_python`
+        fgraph_input_clients = tuple(
+            tuple(
+                (toposort_indices[node], inp_idx)
+                # Disconnect inputs don't have clients
+                for node, inp_idx in clients.get(inp, ())
+            )
+            for inp in fgraph.inputs
+        )
+        fgraph_output_ancestors = tuple(
+            tuple(toposort_indices[inp.owner] for inp in out.owner.inputs)
+            for out in fgraph.outputs
+            if out.owner is not None  # constant outputs
+        )
+
         # Compose individual cache_keys into a global key for the FunctionGraph
         fgraph_key = sha256(
-            f"({type(fgraph)}, {tuple(cache_keys)}, {len(fgraph.inputs)}, {len(fgraph.outputs)})".encode()
+            f"({type(fgraph)}, {tuple(cache_keys)}, {fgraph_input_clients}, {fgraph_output_ancestors})".encode()
         ).hexdigest()
     return numba_njit(py_func), fgraph_key

tests/link/numba/test_basic.py

@@ -7,26 +7,29 @@
 import pytest
 import scipy
 
-from pytensor.compile import SymbolicInput
-from pytensor.tensor.utils import hash_from_ndarray
-
 
 numba = pytest.importorskip("numba")
 
 import pytensor.scalar as ps
 import pytensor.tensor as pt
 from pytensor import config, shared
+from pytensor.compile import SymbolicInput
 from pytensor.compile.function import function
 from pytensor.compile.mode import Mode
 from pytensor.graph.basic import Apply, Variable
+from pytensor.graph.fg import FunctionGraph
 from pytensor.graph.op import Op
 from pytensor.graph.rewriting.db import RewriteDatabaseQuery
 from pytensor.graph.type import Type
 from pytensor.link.numba.dispatch import basic as numba_basic
-from pytensor.link.numba.dispatch.basic import cache_key_for_constant
+from pytensor.link.numba.dispatch.basic import (
+    cache_key_for_constant,
+    numba_funcify_and_cache_key,
+)
 from pytensor.link.numba.linker import NumbaLinker
 from pytensor.scalar.basic import ScalarOp, as_scalar
 from pytensor.tensor.elemwise import Elemwise
+from pytensor.tensor.utils import hash_from_ndarray
 
 
 if TYPE_CHECKING:
@@ -652,3 +655,49 @@ def impl(x):
         outs[2].owner.op, outs[2].owner
     )
     assert numba.njit(lambda x: fn2_def_cached(x))(test_x) == 2
+
+
+def test_fgraph_cache_key():
+    x = pt.scalar("x")
+    log_x = pt.log(x)
+    graphs = [
+        pt.exp(x) / log_x,
+        log_x / pt.exp(x),
+        pt.exp(log_x) / x,
+        x / pt.exp(log_x),
+        pt.exp(log_x) / log_x,
+        log_x / pt.exp(log_x),
+    ]
+
+    def generate_and_validate_key(fg):
+        _, key = numba_funcify_and_cache_key(fg)
+        assert key is not None
+        _, key_again = numba_funcify_and_cache_key(fg)
+        assert key == key_again  # Check its stable
+        return key
+
+    keys = []
+    for graph in graphs:
+        fg = FunctionGraph([x], [graph], clone=False)
+        keys.append(generate_and_validate_key(fg))
+    # Check keys are unique
+    assert len(set(keys)) == len(graphs)
+
+    # Extra unused input should alter the key, because it changes the function signature
+    y = pt.scalar("y")
+    for inputs in [[x, y], [y, x]]:
+        fg = FunctionGraph(inputs, [graphs[0]], clone=False)
+        keys.append(generate_and_validate_key(fg))
+    assert len(set(keys)) == len(graphs) + 2
+
+    # Adding an input as an output should also change the key
+    for outputs in [
+        [graphs[0], x],
+        [x, graphs[0]],
+        [x, x, graphs[0]],
+        [x, graphs[0], x],
+        [graphs[0], x, x],
+    ]:
+        fg = FunctionGraph([x], outputs, clone=False)
+        keys.append(generate_and_validate_key(fg))
+    assert len(set(keys)) == len(graphs) + 2 + 5