torchtyping icon indicating copy to clipboard operation
torchtyping copied to clipboard

Support tensor-likes

Open patrick-kidger opened this issue 3 years ago • 0 comments

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

patrick-kidger avatar May 26 '21 10:05 patrick-kidger