Torchtyping function help strings are extremely verbose

[neelnanda-io]

torchtyping used to give a function type signature is extremely verbose when printing the type, to somewhat ridiculous extents. Eg, here is model.run_with_cache?:

Signature:
model.run_with_cache(
    *model_args,
    return_cache_object=True,
    remove_batch_dim=False,
    **kwargs,
) -> Tuple[Union[NoneType, typing_extensions.Annotated[torch.Tensor, {'__torchtyping__': True, 'details': ('batch', 'pos', 'd_vocab',), 'cls_name': 'TensorType'}], typing_extensions.Annotated[torch.Tensor, {'__torchtyping__': True, 'details': ((),), 'cls_name': 'TensorType'}], typing_extensions.Annotated[torch.Tensor, {'__torchtyping__': True, 'details': ('batch', 'position - 1',), 'cls_name': 'TensorType'}], Tuple[typing_extensions.Annotated[torch.Tensor, {'__torchtyping__': True, 'details': ('batch', 'pos', 'd_vocab',), 'cls_name': 'TensorType'}], Union[typing_extensions.Annotated[torch.Tensor, {'__torchtyping__': True, 'details': ((),), 'cls_name': 'TensorType'}], typing_extensions.Annotated[torch.Tensor, {'__torchtyping__': True, 'details': ('batch', 'position - 1',), 'cls_name': 'TensorType'}]]]], Union[transformer_lens.ActivationCache.ActivationCache, Dict[str, torch.Tensor]]]
Docstring: Wrapper around run_with_cache in HookedRootModule. If return_cache_object is True, this will return an ActivationCache object, with a bunch of useful HookedTransformer specific methods, otherwise it will return a dictionary of activations as in HookedRootModule.

[neelnanda-io]

@alan-cooney thoughts?

[alan-cooney]

Yeah it's a bit of a pain - no obvious solutions unfortunately.

Could be worth creating an issue at https://github.com/patrick-kidger/torchtyping if you're running into this a lot? I'm not sure how common it is to print out type signatures like this

[alan-cooney]

Torchtyping is now encouraging use of https://github.com/google/jaxtyping instead, so this may solve this as well

[jbloomAus]

Seems like this is probably achieved by 164. Can reopen if we notice it's still an issue. I think when formatted nicely, it's probably worth having the typing.