Skip to content

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

@ricardoV94

Description

@ricardoV94

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.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions