from __future__ import annotations

from typing import TYPE_CHECKING

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


# TF type aliases for runtime type *checks* (e.g. with typeguard's check_type function)
TFArray = jxt.Float[tf.Tensor, '...']
TFMatrix = jxt.Float[tf.Tensor, 'Time Features']
TFTensor = jxt.Float[tf.Tensor, 'Batch Time Features']
TFPadMask = jxt.Bool[tf.Tensor, 'Batch Time']
TFSeqLens = jxt.UInt[tf.Tensor, 'Batch']
