Change numba_typify to numba_const_convert

Author: brandonwillard

pytensor/link/numba/dispatch/__init__.py

@@ -1,5 +1,5 @@
 # isort: off
-from pytensor.link.numba.dispatch.basic import numba_funcify, numba_typify
+from pytensor.link.numba.dispatch.basic import numba_funcify, numba_const_convert
 
 # Load dispatch specializations
 import pytensor.link.numba.dispatch.scalar

pytensor/link/numba/dispatch/basic.py

@@ -331,7 +331,8 @@ def use_optimized_cheap_pass(*args, **kwargs):
 
 
 @singledispatch
-def numba_typify(data, dtype=None, **kwargs):
+def numba_const_convert(data, dtype=None, **kwargs):
+    """Create a Numba compatible constant from an PyTensor `Constant`."""
     return data
 
 
@@ -423,7 +424,7 @@ def numba_funcify_FunctionGraph(
     return fgraph_to_python(
         fgraph,
         numba_funcify,
-        const_conversion_fn=numba_typify,
+        const_conversion_fn=numba_const_convert,
         fgraph_name=fgraph_name,
         **kwargs,
     )

pytensor/link/numba/dispatch/random.py

@@ -20,7 +20,7 @@
 from pytensor.graph.basic import Apply
 from pytensor.graph.op import Op
 from pytensor.link.numba.dispatch import basic as numba_basic
-from pytensor.link.numba.dispatch.basic import numba_funcify, numba_typify
+from pytensor.link.numba.dispatch.basic import numba_const_convert, numba_funcify
 from pytensor.link.utils import (
     compile_function_src,
     get_name_for_object,
@@ -96,11 +96,11 @@ def uniform_no_size(a, b, size):
         return uniform_no_size
 
 
-@numba_typify.register(RandomState)
-def numba_typify_RandomState(state, **kwargs):
-    # The numba_typify in this case is just an passthrough function
+@numba_const_convert.register(RandomState)
+def numba_const_convert_RandomState(state, **kwargs):
+    # The `numba_const_convert` in this case is just a passthrough function
     # that synchronizes Numba's internal random state with the current
-    # RandomState object
+    # `RandomState` object.
     ints, index = state.get_state()[1:3]
     ptr = _helperlib.rnd_get_np_state_ptr()
     _helperlib.rnd_set_state(ptr, (index, [int(x) for x in ints]))

pytensor/link/numba/linker.py

@@ -35,13 +35,13 @@ def jit_compile(self, fn):
     def create_thunk_inputs(self, storage_map):
         from numpy.random import RandomState
 
-        from pytensor.link.numba.dispatch import numba_typify
+        from pytensor.link.numba.dispatch import numba_const_convert
 
         thunk_inputs = []
         for n in self.fgraph.inputs:
             sinput = storage_map[n]
             if isinstance(sinput[0], RandomState):
-                new_value = numba_typify(
+                new_value = numba_const_convert(
                     sinput[0], dtype=getattr(sinput[0], "dtype", None)
                 )
                 # We need to remove the reference-based connection to the

tests/link/numba/test_basic.py

@@ -24,7 +24,7 @@
 from pytensor.graph.type import Type
 from pytensor.ifelse import ifelse
 from pytensor.link.numba.dispatch import basic as numba_basic
-from pytensor.link.numba.dispatch import numba_typify
+from pytensor.link.numba.dispatch import numba_const_convert
 from pytensor.link.numba.linker import NumbaLinker
 from pytensor.raise_op import assert_op
 from pytensor.tensor import blas
@@ -321,7 +321,7 @@ def test_create_numba_signature(v, expected, force_scalar):
     [
         (
             np.random.RandomState(1),
-            numba_typify,
+            numba_const_convert,
             lambda x, y: np.all(x.get_state()[1] == y.get_state()[1]),
         )
     ],