Numba linker: fgraph_to_python puts fgraph inputs in global_env, preventing numba caching

Numba linker: fgraph_to_python puts fgraph inputs in global_env, preventing numba caching

Description

When compiling a FunctionGraph with the numba linker, fgraph_to_python
(populates global_env with all non-None storage values, including
variables that are already fgraph.inputs (function parameters). This
means the generated Python function has large array constants in its
__globals__ dict, which numba's CacheImpl.check_cachable rejects as
"dynamic globals (such as ctypes pointers and large global arrays)".

The result: numba cannot cache the compiled artifact, and every fresh
process pays the full JIT cost (~165s for a 30-layer transformer on
Apple M3 Max).

Reproducer

import warnings
import numpy as np
import pytensor
import pytensor.tensor as pt
from pytensor.graph.basic import graph_inputs
from pytensor.compile.sharedvalue import SharedVariable
import pytensor_ml.pytensorf as pytensorf

# Build a graph with shared variable weights
W = pytensor.shared(np.random.randn(64, 64).astype("float32"), name="W")
x = pt.matrix("x")
y = pt.dot(x, W)

# Compile with numba
fn = pytensor.function([x], y, mode="NUMBA")

# Run to trigger JIT
with warnings.catch_warnings(record=True) as w:
    warnings.simplefilter("always")
    fn(np.ones((1, 64), dtype="float32"))

cache_warnings = [x for x in w if "Cannot cache" in str(x.message)]
print(f"Cannot-cache warnings: {len(cache_warnings)}")
for cw in cache_warnings:
    print(f"  {cw.message}")

On pytensor 3.3.2 with numba 0.67.0, this prints:

Cannot-cache warnings: 1
  Cannot cache compiled function "numba_funcified_fgraph" as it uses dynamic globals
  (such as ctypes pointers and large global arrays)

Root cause

In pytensor/link/utils.py, fgraph_to_python:

for node in order:
    for inp in node.inputs:
        is_constant = isinstance(inp, Constant)
        input_storage = storage_map.setdefault(
            inp, [inp.data if isinstance(inp, Constant) else None])
        if (is_constant or input_storage[0] is not None) and inp not in tipifiyed_vars:
            global_env[local_input_name] = type_conversion_fn(
                input_storage[0], variable=inp, storage=input_storage, **kwargs)
            tipifiyed_vars.add(inp)

When a shared variable's value is not None (always true for loaded weights),
it gets placed in global_env even if it's also a fgraph.input. The
generated function then has both the variable as a parameter AND as a global.

For a model like SmolLM2-135M-Instruct (30 layers, 272 weight variables),
this means 332 large array globals (~400MB) in the function's __globals__.

Proposed fix

Skip fgraph inputs from global_env:

# In fgraph_to_python, change:
if (is_constant or input_storage[0] is not None) and inp not in tipifiyed_vars:

# To:
if (is_constant or input_storage[0] is not None) and inp not in tipifiyed_vars and inp not in set(fgraph.inputs):

This is a ~3-line change that makes the generated function only reference
non-input constants in its globals, allowing numba to cache the compiled
artifact.

Impact

For SmolLM2-135M-Instruct on Apple M3 Max:

• Before fix: Every load_llm pays ~165s of numba JIT (5 outer function
  compiles × ~33s each). The 85k entries in ~/.pytensor/numba/ are for
  per-op inner graphs, not the expensive outer functions.
• After fix: The outer functions become cacheable. First process pays
  ~165s; subsequent processes load from numba's disk cache in seconds.

The per-op inner graphs (numba_ofg, 124 compiles, ~44s total) are a
separate, smaller cost that may also benefit from this fix if they have
the same issue.

Environment

• pytensor: 3.3.2
• numba: 0.67.0
• Python: 3.12.13
• macOS 26 / arm64 (Apple M3 Max)

Start in pytensor/link/utils.py at fgraph_to_python and review how node inputs are added to global_env. Run the provided FunctionGraph and numba reproducer, then verify that function inputs are not retained as large globals and that the Cannot cache warning no longer appears.