[Question] Type aliases

[euanong]

Is it possible to define type aliases for commonly used types?

e.g. f["a b c"] for Float[torch.Tensor, "a b c"]

[mivanit]

Here's what I use, no idea if it's best practice. It works with beartype for me, although I've done only pretty limited tests

import typing

import jaxtyping
import torch
import numpy as np

def jaxtype_factory(
		name: str,
		jax_dtype: type,
		array_type: type = jaxtyping.Float,
	) -> type:

	class _BaseArray:
		"""jaxtyping shorthand
		jax_dtype = {jax_dtype}
		array_type = {array_type}
		"""

		__slots__ = ()

		def __new__(cls, *args, **kwargs):
			raise TypeError("Type FArray cannot be instantiated.")

		def __init_subclass__(cls, *args, **kwargs):
			raise TypeError(f"Cannot subclass {cls.__name__}")

		@typing._tp_cache
		def __class_getitem__(cls, params):
			if isinstance(params, str):
				return array_type[jax_dtype, params]
			else:
				raise Exception(f"unexpected type for params:\n{type(params) = }\n{params = }")

	_BaseArray.__name__ = name
	_BaseArray.__doc__ = _BaseArray.__doc__.format(
		jax_dtype=repr(jax_dtype),
		array_type=repr(array_type),
	)

	return _BaseArray

# this makes linters happy
class F_Tensor(torch.Tensor):
	@typing._tp_cache
	def __class_getitem__(cls, params):
		raise NotImplementedError()

F_Tensor = jaxtype_factory("F_Tensor", torch.Tensor, jaxtyping.Float)

[patrick-kidger]

I think the __slots__ is unnecessary on a class that can never be instantiated. :)

But otherwise @mivanit's solution is a good one! I usually recommend writing something like

class f:
    def __class_getitem__(cls, item):
        return Float[torch.Tensor, item]

which is essentially the same thing as @mivanit's solution, but without the bells-and-whistles like docstrings and no-instantiation etc.

Note that this won't work with static type checkers, though. (Support for static type checking is the reason that jaxtyping works the way that it does.) If you do want to support them too then typically the best you can do is something like

from torch import Tensor as loat
from jaxtyping import Float as F

F[loat, "b w h"]

...wihch is obviously a huge hack to work around the heavy limitations of static type checkers.