type-hinting TensorVariable with the OptionalApplyType correctly

[kratsg]

from __future__ import annotations

from typing import Any
import sys

if sys.version_info >= (3, 10):
    from typing import TypeAlias
else:
    from typing_extensions import TypeAlias

from pytensor.graph.basic import Apply
from pytensor.tensor.type import TensorType
from pytensor.tensor.variable import TensorVariable

from pyhs3.typing_compat import NotRequired, TypeAlias

TensorVar: TypeAlias = TensorVariable[TensorType, Apply[Any]]

consider the above (which mypy is happy with). From what I see in pytensor code,

OpType = TypeVar("OpType", bound="Op")
class Apply(Node, Generic[OpType]): ...

which at first glance looks fine, since Apply is a generic type requiring the OpType which is bounded (subtype) of Op. At first, I thought it would be fine to just provide Apply[Op] here, but OptionalApplyType is only compatible with something like Apply[Any] and not Apply[Op]:

OptionalApplyType = TypeVar("OptionalApplyType", None, "Apply", covariant=True)

so it's not clear if I'm treating this wrong or if there's a bug in the code for the typehints here. The main reason I'm defining this TensorVar type, is because I'm building a dictionary mapping names to expressions like this:

import pytensor.tensor as pt

distributions: dict[str, TensorVar] = {}

distributions["sqrt"] = pt.sqrt(pt.scalar('x'))
distributions["sqrt"].eval({'x': 4})

as an example.

[ricardoV94]

Can't you do distributions: dict[str, TensorVariable]? The problem we usually find is that mypy doesn't know the input/output types of Apply nodes coming from some Ops are specifically Tensorvariable subclasses of Variable

[kratsg]

I am doing that, but I want to type-hint this strictly, and this requires not leaving generic types around (explicitly typing with Any or similar).

[ricardoV94]

Re:

but OptionalApplyType is only compatible with something like Apply[Any] and not Apply[Op].

I don't know. The point is that a Variable can have an owner (an Apply[Op]) or no owner (None). Not sure we are doing something wrong?