Allow lazy import of linker dispatchers

Author: ricardoV94

pytensor/link/jax/dispatch/__init__.py

@@ -1,5 +1,5 @@
 # isort: off
-from pytensor.link.jax.dispatch.basic import jax_funcify, jax_typify
+from pytensor.link.jax.linker import jax_funcify, jax_typify
 
 # Load dispatch specializations
 import pytensor.link.jax.dispatch.blas

pytensor/link/jax/dispatch/basic.py

@@ -1,6 +1,5 @@
 import warnings
 from collections.abc import Callable
-from functools import singledispatch
 
 import jax
 import jax.numpy as jnp
@@ -10,8 +9,10 @@
 from pytensor.compile.builders import OpFromGraph
 from pytensor.compile.ops import DeepCopyOp, ViewOp
 from pytensor.configdefaults import config
+from pytensor.graph import Op
 from pytensor.graph.fg import FunctionGraph
 from pytensor.ifelse import IfElse
+from pytensor.link.jax.linker import jax_funcify, jax_typify
 from pytensor.link.utils import fgraph_to_python
 from pytensor.raise_op import Assert, CheckAndRaise
 
@@ -22,24 +23,15 @@
     jax.config.update("jax_enable_x64", False)
 
 
-@singledispatch
-def jax_typify(data, dtype=None, **kwargs):
-    r"""Convert instances of PyTensor `Type`\s to JAX types."""
-    if dtype is None:
-        return data
-    else:
-        return jnp.array(data, dtype=dtype)
-
-
 @jax_typify.register(np.ndarray)
 def jax_typify_ndarray(data, dtype=None, **kwargs):
     if len(data.shape) == 0:
         return data.item()
     return jnp.array(data, dtype=dtype)
 
 
-@singledispatch
-def jax_funcify(op, node=None, storage_map=None, **kwargs):
+@jax_funcify.register(Op)
+def jax_funcify_op(op, node=None, storage_map=None, **kwargs):
     """Create a JAX compatible function from an PyTensor `Op`."""
     raise NotImplementedError(f"No JAX conversion for the given `Op`: {op}")
 

pytensor/link/jax/linker.py

@@ -1,11 +1,29 @@
 import warnings
+from functools import singledispatch
 
 from numpy.random import Generator
 
 from pytensor.compile.sharedvalue import SharedVariable, shared
 from pytensor.link.basic import JITLinker
 
 
+@singledispatch
+def jax_typify(data, dtype=None, **kwargs):
+    r"""Convert instances of PyTensor `Type`\s to JAX types."""
+    import jax.numpy as jnp
+
+    if dtype is None:
+        return data
+    else:
+        return jnp.array(data, dtype=dtype)
+
+
+@singledispatch
+def jax_funcify(obj, *args, **kwargs):
+    """Create a JAX compatible function from an PyTensor `Op`."""
+    raise NotImplementedError(f"No JAX conversion for the given type: {type(obj)}")
+
+
 class JAXLinker(JITLinker):
     """A `Linker` that JIT-compiles NumPy-based operations using JAX."""
 
@@ -14,7 +32,6 @@ def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
 
     def fgraph_convert(self, fgraph, input_storage, storage_map, **kwargs):
-        from pytensor.link.jax.dispatch import jax_funcify
         from pytensor.link.jax.dispatch.shape import JAXShapeTuple
         from pytensor.tensor.random.type import RandomType
 
@@ -111,8 +128,6 @@ def convert_scalar_shape_inputs(
         return convert_scalar_shape_inputs
 
     def create_thunk_inputs(self, storage_map):
-        from pytensor.link.jax.dispatch import jax_typify
-
         thunk_inputs = []
         for n in self.fgraph.inputs:
             sinput = storage_map[n]

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.linker import numba_funcify, numba_typify
 
 # Load dispatch specializations
 import pytensor.link.numba.dispatch.blockwise

pytensor/link/numba/dispatch/basic.py

@@ -2,7 +2,6 @@
 import sys
 import warnings
 from copy import copy
-from functools import singledispatch
 from textwrap import dedent
 
 import numba
@@ -21,11 +20,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 Op
 from pytensor.graph.basic import Apply
 from pytensor.graph.fg import FunctionGraph
 from pytensor.graph.type import Type
 from pytensor.ifelse import IfElse
 from pytensor.link.numba.dispatch.sparse import CSCMatrixType, CSRMatrixType
+from pytensor.link.numba.linker import numba_funcify, numba_typify
 from pytensor.link.utils import (
     compile_function_src,
     fgraph_to_python,
@@ -276,11 +277,6 @@ def create_arg_string(x):
     return args
 
 
-@singledispatch
-def numba_typify(data, dtype=None, **kwargs):
-    return data
-
-
 def generate_fallback_impl(op, node=None, storage_map=None, **kwargs):
     """Create a Numba compatible function from a Pytensor `Op`."""
 
@@ -326,14 +322,8 @@ def perform(*inputs):
     return perform
 
 
-@singledispatch
-def numba_funcify(op, node=None, storage_map=None, **kwargs):
-    """Generate a numba function for a given op and apply node.
-
-    The resulting function will usually use the `no_cpython_wrapper`
-    argument in numba, so it can not be called directly from python,
-    but only from other jit functions.
-    """
+@numba_funcify.register(Op)
+def numba_funcify_op(op, node=None, storage_map=None, **kwargs):
     return generate_fallback_impl(op, node, storage_map, **kwargs)
 
 

pytensor/link/numba/linker.py

@@ -1,12 +1,24 @@
+from functools import singledispatch
+
 from pytensor.link.basic import JITLinker
 
 
+@singledispatch
+def numba_typify(data, dtype=None, **kwargs):
+    raise NotImplementedError(
+        f"Numba funcify not implemented for data type {type(data)}"
+    )
+
+
+@singledispatch
+def numba_funcify(obj, *args, **kwargs):
+    raise NotImplementedError(f"Numba funcify not implemented for type {type(obj)}")
+
+
 class NumbaLinker(JITLinker):
     """A `Linker` that JIT-compiles NumPy-based operations using Numba."""
 
     def fgraph_convert(self, fgraph, **kwargs):
-        from pytensor.link.numba.dispatch import numba_funcify
-
         return numba_funcify(fgraph, **kwargs)
 
     def jit_compile(self, fn):

pytensor/link/pytorch/dispatch/__init__.py

@@ -1,5 +1,5 @@
 # isort: off
-from pytensor.link.pytorch.dispatch.basic import pytorch_funcify, pytorch_typify
+from pytensor.link.pytorch.linker import pytorch_funcify, pytorch_typify
 
 # # Load dispatch specializations
 import pytensor.link.pytorch.dispatch.blas

pytensor/link/pytorch/dispatch/basic.py

@@ -1,4 +1,3 @@
-from functools import singledispatch
 from types import NoneType
 
 import numpy as np
@@ -10,9 +9,11 @@
 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 Op
 from pytensor.graph.basic import Constant
 from pytensor.graph.fg import FunctionGraph
 from pytensor.ifelse import IfElse
+from pytensor.link.pytorch.linker import pytorch_funcify, pytorch_typify
 from pytensor.link.utils import fgraph_to_python
 from pytensor.raise_op import CheckAndRaise
 from pytensor.tensor.basic import (
@@ -27,11 +28,6 @@
 )
 
 
-@singledispatch
-def pytorch_typify(data, **kwargs):
-    raise NotImplementedError(f"pytorch_typify is not implemented for {type(data)}")
-
-
 @pytorch_typify.register(np.ndarray)
 @pytorch_typify.register(torch.Tensor)
 def pytorch_typify_tensor(data, dtype=None, **kwargs):
@@ -45,8 +41,8 @@ def pytorch_typify_no_conversion_needed(data, **kwargs):
     return data
 
 
-@singledispatch
-def pytorch_funcify(op, node=None, storage_map=None, **kwargs):
+@pytorch_funcify.register(Op)
+def pytorch_funcify_op(op, node=None, storage_map=None, **kwargs):
     """Create a PyTorch compatible function from an PyTensor `Op`."""
     raise NotImplementedError(
         f"No PyTorch conversion for the given `Op`: {op}.\nCheck out `https://github.com/pymc-devs/pytensor/issues/821` for progress or to request we prioritize this operation"

pytensor/link/pytorch/linker.py

@@ -1,7 +1,19 @@
+from functools import singledispatch
+
 from pytensor.link.basic import JITLinker
 from pytensor.link.utils import unique_name_generator
 
 
+@singledispatch
+def pytorch_typify(data, **kwargs):
+    raise NotImplementedError(f"pytorch_typify is not implemented for {type(data)}")
+
+
+@singledispatch
+def pytorch_funcify(obj, *args, **kwargs):
+    raise NotImplementedError(f"pytorch_funcify is not implemented for {type(obj)}")
+
+
 class PytorchLinker(JITLinker):
     """A `Linker` that compiles NumPy-based operations using torch.compile."""
 
@@ -10,8 +22,6 @@ def __init__(self, *args, **kwargs):
         self.gen_functors = []
 
     def fgraph_convert(self, fgraph, input_storage, storage_map, **kwargs):
-        from pytensor.link.pytorch.dispatch import pytorch_funcify
-
         # We want to have globally unique names
         # across the entire pytensor graph, not
         # just the subgraph
@@ -40,8 +50,6 @@ def jit_compile(self, fn):
         # flag that tend to help our graphs
         torch._dynamo.config.capture_dynamic_output_shape_ops = True
 
-        from pytensor.link.pytorch.dispatch import pytorch_typify
-
         class wrapper:
             """
             Pytorch would fail compiling our method when trying