from __future__ import annotations

from typing import TYPE_CHECKING

import jaxtyping as jxt
try:
    import torch
except ImportError:
    if TYPE_CHECKING:
        from . import mock_types as torch  # type: ignore[no-redef]
    else:
        raise ImportError(
            'torch not installed - try one of the following:\n'
            '-> with pip: pip install tempora-ai[torch]\n'
            '-> with poetry: poetry add tempora-ai -E torch',
        )


# PyTorch type aliases for runtime type *checks* (e.g. with typeguard's check_type func)
TorchArray = jxt.Float[torch.Tensor, '...']
TorchMatrix = jxt.Float[torch.Tensor, 'Time Features']
TorchTensor = jxt.Float[torch.Tensor, 'Batch Time Features']
TorchPadMask = jxt.Bool[torch.Tensor, 'Batch Time']
TorchSeqLens = jxt.Int[torch.Tensor, 'Batch']  # Torch only supports 8-bit unsigned ints
