torchtyping
torchtyping copied to clipboard
Support tensor-likes
That is, classes supporting the __torch_function__
protocol.
This shouldn't be too difficult -- most of the necessary work has already been done.
- There's some places where we have
torch.Tensor
hardcoded, for example in instance checks and some type annotations, that would need adjusting to accept tensor-likes. -
TensorTypeMixin
would need exposing as a public part of the interface. - The documentation needs updating to show how this is possible. Once the above changes are made it should just be:
from torchtyping import TensorTypeMixin
class TensorLike:
...
class TensorLikeType(TensorLike, TensorTypeMixin):
base_cls = TensorLike