[Question] Type aliases about jaxtyping HOT 2 OPEN

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

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

F[loat, "b w h"]

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)