| |
|
|
| """ |
| This module contains type annotations for the project, using |
| 1. Python type hints (https://docs.python.org/3/library/typing.html) for Python objects |
| 2. jaxtyping (https://github.com/google/jaxtyping/blob/main/API.md) for PyTorch tensors |
| |
| Two types of typing checking can be used: |
| 1. Static type checking with mypy (install with pip and enabled as the default linter in VSCode) |
| 2. Runtime type checking with typeguard (install with pip and triggered at runtime, mainly for tensor dtype and shape checking) |
| """ |
|
|
| |
| from typing import ( |
| Any, |
| Callable, |
| Dict, |
| Iterable, |
| List, |
| Literal, |
| NamedTuple, |
| NewType, |
| Optional, |
| Sized, |
| Tuple, |
| Type, |
| TypeVar, |
| Union, |
| ) |
|
|
| |
| |
| from jaxtyping import Bool, Complex, Float, Inexact, Int, Integer, Num, Shaped, UInt |
|
|
| |
| from omegaconf import DictConfig |
|
|
| |
| from torch import Tensor |
|
|
| |
| from typeguard import typechecked as typechecker |
|
|