pymc-devs / pymc-devs/pytensor
Add a type converter registry to Type to replace cross-package deferred imports
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 644
- Forks
- 208
- Avg merge
- 2d 14h
- Merged PRs (30d)
- 16
Description
Problem
Several Type subclasses use deferred imports in convert_variable (and similar methods) to handle conversions across package boundaries:
scalar → tensor (pytensor/scalar/basic.py):
# ScalarType reaching up to tensor
from pytensor.tensor.basic import as_tensor_variable
from pytensor.tensor.basic import scalar_from_tensor
from pytensor.tensor.type import TensorType
tensor → xtensor (pytensor/tensor/random/type.py, introduced in the rng_methods PR):
# RandomGeneratorType reaching across to xtensor
from pytensor.xtensor.random.type import XRandomGeneratorType, xrng_to_rng
In all cases the dependency direction is inverted — lower-level packages reach into higher-level ones via lazy imports. It works, but each conversion must be hardcoded in the source type's package, and adding a new type system layer means editing existing code.
Proposal
A registration pattern on Type itself, similar to shared_constructor.register:
# In pytensor/graph/type.py
class Type:
_type_converters: ClassVar[dict[type, Callable]] = {}
@classmethod
def register_converter(cls, from_type, converter_fn):
cls._type_converters[from_type] = converter_fn
def convert_variable(self, var):
converter = self._type_converters.get(type(var.type))
if converter is not None:
return converter(var)
# existing is_super logic...
Each downstream package registers its converters at import time:
# xtensor/random/type.py registers on import
RandomGeneratorType.register_converter(XRandomGeneratorType, xrng_to_rng)
# tensor/basic.py registers on import
ScalarType.register_converter(TensorType, scalar_from_tensor)
This restores the correct dependency direction — higher-level packages extend lower-level types, not the other way around.
Scope
Refactor with no user-facing behavior change. Currently affects:
ScalarType.convert_variable(scalar↔tensor)RandomGeneratorType.convert_variable(tensor↔xtensor RNG)as_rngintensor/random/type.py
Could also clean up any future xtensor↔tensor conversions as xtensor matures.
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start in pytensor/graph/type.py, then inspect ScalarType.convert_variable in pytensor/scalar/basic.py, RandomGeneratorType.convert_variable and as_rng in pytensor/tensor/random/type.py, and the related registrations in pytensor/tensor/basic.py and pytensor/xtensor/random/type.py. Trace the existing conversion paths before choosing how registration and import-time setup should work. Done means the deferred cross-package imports are replaced without changing conversion behavior.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- backend
- Issue type
- Refactor
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 48/100