Make Op class an immutable Generic#1997
Conversation
pytensor/scan/op.py
Outdated
|
|
||
|
|
||
| class Scan(Op, ScanMethodsMixin, HasInnerGraph): | ||
| class Scan(Op[TensorVariable], ScanMethodsMixin, HasInnerGraph): |
There was a problem hiding this comment.
Scan can output any variable type (as untraced_sit_sot)
The Scan materialized with an inner graph should determine the output type
pytensor/sparse/basic.py
Outdated
|
|
||
|
|
||
| class CSM(Op): | ||
| class CSM(Op[TensorVariable]): |
There was a problem hiding this comment.
Should be SparseTensorVariable?
There was a problem hiding this comment.
Is there a such a thing? It would be great if there were, but I didn't find it.
There was a problem hiding this comment.
There is,
pytensor/pytensor/sparse/variable.py
Line 429 in c311830
SparseVariable not SparseTensorVariable. Making a subclass of TensorVariable was still a mistake imo, so we should change that later down the road
pytensor/tensor/random/op.py
Outdated
|
|
||
|
|
||
| class RNGConsumerOp(Op): | ||
| class RNGConsumerOp(Op[TensorVariable]): |
There was a problem hiding this comment.
also outputs RandomType variable (the next rng state)
There was a problem hiding this comment.
I don't know how to type hint that 😬😬😬
My type hints only work for ops that always return the same type. I don't know how to add more output types on the fly.
There was a problem hiding this comment.
The output signature of RandomVariable (not RNGConsumerOp) is tuple[RandomTypeVariable, TensorVariable]. However there is the annoying default_output, which means the signature of call is just Tensorvariable. What about multi-output ops in general?
| raise TypeError( | ||
| f"The 'outputs' argument to Apply must contain Variable instances with no owner, not {output}" | ||
| ) | ||
| self.outputs = cast(ApplyOutputsType, tuple(_outputs)) |
There was a problem hiding this comment.
Sounds nit, but apply creation is hot loop for us, don't wast cycles with cast function (even a no-op like this), just use type-ignore if mypy doesn't believe us
2 participants
This PR aims to solve a subset of the type hint issues that
pytensorhas, namely theOpclass expects to haveVariableinput instances and outputVariableinstances. This makes it impossible for type checkers to realize that anOpactually operates withVariablesubtypes likeTensorVariable, and so the type checker will never know that an Op's output will have all of the tensor traits.The PR tries to handle this by making
OpaGenericclass. After much thought, I decided to make theOpan immutableGeneric. I decided against covariant or contravariant because subclasses of subclasses ofOp's might have outputs that are not subclasses of their parent class output, so immutable seemed like the safest choice.