pymc-devs / pymc-devs/pytensor

Add a type converter registry to Type to replace cross-package deferred imports

Open
#2,029 0 comments 0 reactions 0 assignees View on GitHub

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_rng in tensor/random/type.py

Could also clean up any future xtensor↔tensor conversions as xtensor matures.

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.