nshtrainer 0.1.0__py3-none-any.whl

nshtrainer/trainer/signal_connector.py

@@ -0,0 +1,208 @@
+import logging
+import os
+import re
+import signal
+import subprocess
+import threading
+from collections import defaultdict
+from collections.abc import Callable
+from pathlib import Path
+from types import FrameType
+from typing import Any, TypeAlias
+
+import torch.utils.data
+from lightning.fabric.plugins.environments.lsf import LSFEnvironment
+from lightning.fabric.plugins.environments.slurm import SLURMEnvironment
+from lightning.pytorch.trainer.connectors.signal_connector import _HandlersCompose
+from lightning.pytorch.trainer.connectors.signal_connector import (
+    _SignalConnector as _LightningSignalConnector,
+)
+from typing_extensions import override
+
+log = logging.getLogger(__name__)
+
+_SIGNUM = int | signal.Signals
+_HANDLER: TypeAlias = Callable[[_SIGNUM, FrameType], Any] | int | signal.Handlers | None
+
+
+class _SignalConnector(_LightningSignalConnector):
+    def _auto_requeue_signals(self) -> list[signal.Signals]:
+        from ..model.base import BaseConfig
+
+        if not isinstance(config := self.trainer.lightning_module.hparams, BaseConfig):
+            return []
+
+        signals = config.runner.submit._resolved_auto_requeue_signals()
+        signals_set = set(signals)
+        valid_signals: set[signal.Signals] = signal.valid_signals()
+        assert signals_set.issubset(
+            valid_signals
+        ), f"Invalid signal(s) found: {signals_set - valid_signals}"
+        return signals
+
+    def _compose_and_register(
+        self,
+        signum: _SIGNUM,
+        handlers: list[_HANDLER],
+        replace_existing: bool = False,
+    ):
+        if self._is_on_windows():
+            log.info(f"Signal {signum} has no handlers or is not supported on Windows.")
+            return
+
+        if self._has_already_handler(signum):
+            if not replace_existing:
+                log.info(
+                    f"Signal {signum} already has a handler. Adding ours to the existing one."
+                )
+                handlers.append(signal.getsignal(signum))
+            else:
+                log.info(f"Replacing existing handler for signal {signum} with ours.")
+
+        self._register_signal(signum, _HandlersCompose(handlers))
+        log.info(f"Registered {len(handlers)} handlers for signal {signum}.")
+
+    @override
+    def register_signal_handlers(self) -> None:
+        if not (auto_requeue_signals := self._auto_requeue_signals()):
+            log.info(
+                "No auto-requeue signals found. Reverting to default Lightning behavior."
+            )
+            return super().register_signal_handlers()
+
+        self.received_sigterm = False
+        self._original_handlers = self._get_current_signal_handlers()
+
+        signals = defaultdict[signal.Signals, list[_HANDLER]](lambda: [])
+        signals[signal.SIGTERM].append(self._sigterm_notifier_fn)
+
+        environment = self.trainer._accelerator_connector.cluster_environment
+        if isinstance(environment, SLURMEnvironment):
+            log.info("SLURM auto-requeueing enabled. Setting signal handlers.")
+            for signal_handler in auto_requeue_signals:
+                signals[signal_handler].append(self._slurm_sigusr_handler_fn)
+
+        if isinstance(environment, LSFEnvironment):
+            # Important note from https://amrex-astro.github.io/workflow/olcf-workflow.html:
+            # We can also ask the job manager to send a warning signal some amount of time before the allocation expires by passing -wa 'signal' and -wt '[hour:]minute' to bsub. We can then have bash create a dump_and_stop file when it receives the signal, which will tell Castro to output a checkpoint file and exit cleanly after it finishes the current timestep. An important detail that I couldn't find documented anywhere is that the job manager sends the signal to all the processes in the job, not just the submission script, and we have to use a signal that is ignored by default so Castro doesn't immediately crash upon receiving it. SIGCHLD, SIGURG, and SIGWINCH are the only signals that fit this requirement and of these, SIGURG is the least likely to be triggered by other events.
+
+            log.info("LSF auto-requeueing enabled. Setting signal handlers.")
+            for signal_handler in auto_requeue_signals:
+                signals[signal_handler].append(self._lsf_sigusr_handler_fn)
+
+        for signum, handlers in signals.items():
+            if not handlers:
+                continue
+
+            self._compose_and_register(signum, handlers)
+
+    def _should_ignore_signal_handler(self) -> str | None:
+        if threading.current_thread() is not threading.main_thread():
+            return "Not in main thread"
+
+        if torch.utils.data.get_worker_info() is not None:
+            return "Inside DataLoader worker process"
+
+        return None
+
+    @override
+    def _slurm_sigusr_handler_fn(self, signum: _SIGNUM, _: FrameType) -> None:
+        if ignore_reason := self._should_ignore_signal_handler():
+            log.info(
+                f"Skipping SLURM auto-requeue signal handler. Reason: {ignore_reason}"
+            )
+            return
+
+        log.critical(f"Handling SLURM auto-requeue signal: {signum}")
+
+        # save logger to make sure we get all the metrics
+        for logger in self.trainer.loggers:
+            logger.finalize("finished")
+
+        hpc_save_path = self.trainer._checkpoint_connector.hpc_save_path(
+            self.trainer.default_root_dir
+        )
+        self.trainer.save_checkpoint(hpc_save_path)
+
+        if not self.trainer.is_global_zero:
+            return
+
+        # find job id
+        array_job_id = os.getenv("SLURM_ARRAY_JOB_ID")
+        if array_job_id is not None:
+            array_task_id = os.environ["SLURM_ARRAY_TASK_ID"]
+            job_id = f"{array_job_id}_{array_task_id}"
+        else:
+            job_id = os.environ["SLURM_JOB_ID"]
+
+        assert re.match("[0-9_-]+", job_id)
+        cmd = ["scontrol", "requeue", job_id]
+
+        # requeue job
+        log.info(f"requeing job {job_id}...")
+        try:
+            result = subprocess.call(cmd)
+        except FileNotFoundError:
+            # This can occur if a subprocess call to `scontrol` is run outside a shell context
+            # Re-attempt call (now with shell context). If any error is raised, propagate to user.
+            # When running a shell command, it should be passed as a single string.
+            result = subprocess.call(" ".join(cmd), shell=True)
+
+        # print result text
+        if result == 0:
+            log.info(f"Requeued SLURM job: {job_id}")
+        else:
+            log.warning(f"Requeuing SLURM job {job_id} failed with error code {result}")
+
+    def _lsf_sigusr_handler_fn(self, signum: _SIGNUM, _: FrameType) -> None:
+        if ignore_reason := self._should_ignore_signal_handler():
+            log.info(
+                f"Skipping LSF auto-requeue signal handler. Reason: {ignore_reason}"
+            )
+            return
+
+        log.critical(f"Handling LSF auto-requeue signal: {signum}")
+
+        # Save logger to make sure we get all the metrics
+        for logger in self.trainer.loggers:
+            logger.finalize("finished")
+
+        # Save checkpoint
+        hpc_save_path = self.trainer._checkpoint_connector.hpc_save_path(
+            self.trainer.default_root_dir
+        )
+        self.trainer.save_checkpoint(hpc_save_path)
+        log.info(f"Saved checkpoint to {hpc_save_path}")
+
+        if not self.trainer.is_global_zero:
+            return
+
+        # Find job id
+        if (job_id := os.getenv("LSB_JOBID")) is None:
+            log.warning(
+                "LSB_JOBID environment variable not found. Unable to requeue job."
+            )
+            return
+
+        assert re.match("[0-9_-]+", job_id)
+
+        exe = "brequeue"
+        if (bin_dir := os.getenv("LSF_BINDIR")) is not None:
+            exe = str((Path(bin_dir) / exe).resolve().absolute())
+
+        log.info(f"Using LSF requeue executable: {exe}")
+        cmd = [exe, job_id]
+
+        # Requeue job
+        log.info(f"Requeuing job {job_id}...")
+        try:
+            result = subprocess.call(cmd)
+        except FileNotFoundError:
+            # Retry with shell context if subprocess call fails
+            result = subprocess.call(" ".join(cmd), shell=True)
+
+        # Print result text
+        if result == 0:
+            log.info(f"Requeued LSF job: {job_id}")
+        else:
+            log.warning(f"Requeuing LSF job {job_id} failed with error code {result}")

nshtrainer/trainer/trainer.py

@@ -0,0 +1,340 @@
+import contextlib
+import logging
+import os
+from collections.abc import Sequence
+from pathlib import Path
+from typing import Any, cast
+
+import torch
+from lightning.fabric.plugins.environments.lsf import LSFEnvironment
+from lightning.fabric.plugins.environments.slurm import SLURMEnvironment
+from lightning.fabric.plugins.precision.precision import _PRECISION_INPUT
+from lightning.pytorch import LightningModule
+from lightning.pytorch import Trainer as LightningTrainer
+from lightning.pytorch.profilers import Profiler
+from lightning.pytorch.utilities.types import _EVALUATE_OUTPUT, _PREDICT_OUTPUT
+from typing_extensions import Unpack, assert_never, override
+
+from ..actsave import ActSave
+from ..callbacks.base import resolve_all_callbacks
+from ..model.config import (
+    AcceleratorConfigProtocol,
+    BaseConfig,
+    BaseProfilerConfig,
+    LightningTrainerKwargs,
+    StrategyConfigProtocol,
+)
+from .signal_connector import _SignalConnector
+
+log = logging.getLogger(__name__)
+
+
+def _is_bf16_supported_no_emulation():
+    r"""Return a bool indicating if the current CUDA/ROCm device supports dtype bfloat16."""
+    version = cast(Any, torch.version)
+
+    # Check for ROCm, if true return true, no ROCM_VERSION check required,
+    # since it is supported on AMD GPU archs.
+    if version.hip:
+        return True
+
+    device = torch.cuda.current_device()
+
+    # Check for CUDA version and device compute capability.
+    # This is a fast way to check for it.
+    cuda_version = version.cuda
+    if (
+        cuda_version is not None
+        and int(cuda_version.split(".")[0]) >= 11
+        and torch.cuda.get_device_properties(device).major >= 8
+    ):
+        return True
+
+    return False
+
+
+class Trainer(LightningTrainer):
+    @classmethod
+    @contextlib.contextmanager
+    def context(cls, config: BaseConfig):
+        if (precision := config.trainer.set_float32_matmul_precision) is not None:
+            torch.set_float32_matmul_precision(precision)
+
+        yield
+
+    @classmethod
+    def _update_kwargs(
+        cls,
+        config: BaseConfig,
+        kwargs_ctor: LightningTrainerKwargs,
+    ):
+        kwargs: LightningTrainerKwargs = {
+            "deterministic": config.trainer.reproducibility.deterministic,
+            "fast_dev_run": config.trainer.fast_dev_run,
+            "max_epochs": config.trainer.max_epochs,
+            "min_epochs": config.trainer.min_epochs,
+            "max_steps": config.trainer.max_steps,
+            "min_steps": config.trainer.min_steps,
+            "max_time": config.trainer.max_time,
+            "limit_train_batches": config.trainer.limit_train_batches,
+            "limit_val_batches": config.trainer.limit_val_batches,
+            "limit_test_batches": config.trainer.limit_test_batches,
+            "limit_predict_batches": config.trainer.limit_predict_batches,
+            "overfit_batches": config.trainer.overfit_batches,
+            "val_check_interval": config.trainer.val_check_interval,
+            "num_sanity_val_steps": config.trainer.num_sanity_val_steps,
+            "log_every_n_steps": config.trainer.log_every_n_steps,
+            "inference_mode": config.trainer.inference_mode,
+            "callbacks": [],
+            "plugins": [],
+            "logger": [],
+            # Moved to `lightning_kwargs`:
+            # "enable_checkpointing": config.trainer.enable_checkpointing,
+            # "accelerator": config.trainer.accelerator,
+            # "strategy": config.trainer.strategy,
+            # "num_nodes": config.trainer.num_nodes,
+            # "precision": config.trainer.precision,
+            # "logger": config.trainer.logging.enabled,
+            # "log_every_n_steps": config.trainer.log_every_n_steps,
+            # "enable_progress_bar": config.trainer.enable_progress_bar,
+            # "enable_model_summary": config.trainer.enable_model_summary,
+            # "accumulate_grad_batches": config.trainer.accumulate_grad_batches,
+            # "benchmark": config.trainer.benchmark,
+            # "use_distributed_sampler": config.trainer.use_distributed_sampler,
+            # "detect_anomaly": config.trainer.detect_anomaly,
+            # "barebones": config.trainer.barebones,
+            # "plugins": config.trainer.plugins,
+            # "sync_batchnorm": config.trainer.sync_batchnorm,
+            # "reload_dataloaders_every_n_epochs": config.trainer.reload_dataloaders_every_n_epochs,
+        }
+
+        def _update_key(key: str, new_value: Any):
+            # First, check to see if the key is already in the kwargs.
+            if key not in kwargs:
+                kwargs[key] = new_value
+                return
+
+            # If the key is already in the kwargs, then we check the type:
+            # - If the type is a sequence, then we extend the sequence.
+            # - Otherwise, we just update the value but warn the user.
+
+            match existing_value := kwargs[key]:
+                case Sequence() as existing_value:
+                    # Make sure value is a sequence too
+                    if not isinstance(new_value, Sequence):
+                        new_value = [new_value]
+                    kwargs[key] = [*existing_value, *new_value]
+                case _:
+                    log.warning(
+                        f"Trainer.__init__: Overwriting existing value {existing_value=} with {new_value=} for key {key=}."
+                    )
+                    kwargs[key] = new_value
+
+        def _update_kwargs(**update: Unpack[LightningTrainerKwargs]):
+            for key, value in update.items():
+                _update_key(key, value)
+
+        # Set `default_root_dir` if `auto_set_default_root_dir` is enabled.
+        if config.trainer.auto_set_default_root_dir:
+            if kwargs.get("default_root_dir"):
+                raise ValueError(
+                    "You have set `config.trainer.default_root_dir`. "
+                    "But we are trying to set it automatically. "
+                    "Please use `config.directory.base` rather than `config.trainer.default_root_dir`. "
+                    "If you want to set it manually, please set `config.trainer.auto_set_default_root_dir=False`."
+                )
+
+            _update_kwargs(
+                default_root_dir=config.directory.resolve_run_root_directory(config.id)
+            )
+
+        if (devices_input := config.trainer.devices) is not None:
+            match devices_input:
+                case "all":
+                    devices = -1
+                case "auto":
+                    devices = "auto"
+                case Sequence():
+                    devices = list(devices_input)
+                case _:
+                    raise ValueError(f"Invalid value for devices={devices_input}.")
+
+            _update_kwargs(devices=devices)
+
+        if (
+            use_distributed_sampler := config.trainer.use_distributed_sampler
+        ) is not None:
+            _update_kwargs(use_distributed_sampler=use_distributed_sampler)
+
+        if (accelerator := config.trainer.accelerator) is not None:
+            if isinstance(accelerator, AcceleratorConfigProtocol):
+                accelerator = accelerator.construct_accelerator()
+            _update_kwargs(accelerator=accelerator)
+
+        if (strategy := config.trainer.strategy) is not None:
+            if isinstance(strategy, StrategyConfigProtocol):
+                strategy = strategy.construct_strategy()
+            _update_kwargs(strategy=strategy)
+
+        if (precision := config.trainer.precision) is not None:
+            resolved_precision: _PRECISION_INPUT
+            match precision:
+                case "64-true" | "32-true" | "bf16-mixed":
+                    resolved_precision = precision
+                case "fp16-mixed":
+                    resolved_precision = "16-mixed"
+                case "16-mixed-auto":
+                    try:
+                        resolved_precision = (
+                            "bf16-mixed"
+                            if _is_bf16_supported_no_emulation()
+                            else "16-mixed"
+                        )
+                    except BaseException:
+                        resolved_precision = "16-mixed"
+                        log.warning(
+                            "Failed to detect bfloat16 support. Falling back to 16-mixed."
+                        )
+
+                    log.critical(
+                        f"Auto-resolving {precision=} to {resolved_precision=}."
+                    )
+                case _:
+                    assert_never(precision)
+
+            _update_kwargs(precision=resolved_precision)
+
+        if (detect_anomaly := config.trainer.detect_anomaly) is not None:
+            _update_kwargs(detect_anomaly=detect_anomaly)
+
+        if (
+            grad_clip_config := config.trainer.optimizer.gradient_clipping
+        ) is not None and grad_clip_config.enabled:
+            # kwargs["gradient_clip_algorithm"] = grad_clip_config.algorithm
+            # kwargs["gradient_clip_val"] = grad_clip_config.value
+            _update_kwargs(
+                gradient_clip_algorithm=grad_clip_config.algorithm,
+                gradient_clip_val=grad_clip_config.value,
+            )
+
+        if profiler := config.trainer.profiler:
+            # If the profiler is an ProfilerConfig instance, then we instantiate it.
+            if isinstance(profiler, BaseProfilerConfig):
+                profiler = profiler.construct_profiler(config)
+                # Make sure that the profiler is an instance of `Profiler`.
+                if not isinstance(profiler, Profiler):
+                    raise ValueError(f"{profiler=} is not an instance of `{Profiler}`.")
+
+            # Otherwise, if the profiler is a string (e.g., "simpe", "advanced", "pytorch"),
+            #   then we just pass it through.
+            # kwargs["profiler"] = profiler
+            _update_kwargs(profiler=profiler)
+
+        if callbacks := resolve_all_callbacks(config):
+            _update_kwargs(callbacks=callbacks)
+
+        if plugin_configs := config.trainer.plugins:
+            _update_kwargs(
+                plugins=[
+                    plugin_config.construct_plugin() for plugin_config in plugin_configs
+                ]
+            )
+
+        if not config.trainer.logging.enabled:
+            log.critical(f"Disabling logger because {config.trainer.logging.enabled=}.")
+            kwargs["logger"] = False
+        else:
+            _update_kwargs(logger=config.trainer.logging.construct_loggers(config))
+
+        if config.trainer.auto_determine_num_nodes:
+            # When num_nodes is auto, we need to detect the number of nodes.
+            if SLURMEnvironment.detect():
+                if (num_nodes := os.environ.get("SLURM_NNODES")) is not None:
+                    num_nodes = int(num_nodes)
+                    log.critical(f"SLURM detected with {num_nodes=}.")
+                    _update_kwargs(num_nodes=num_nodes)
+                else:
+                    log.critical(
+                        "SLURM detected, but SLURM_NNODES not found. "
+                        "We'll continue without setting num_nodes, but this may cause issues."
+                    )
+
+            elif LSFEnvironment.detect():
+                num_nodes = LSFEnvironment().world_size()
+                log.critical(f"LSF detected with {num_nodes=}.")
+                _update_kwargs(num_nodes=num_nodes)
+            else:
+                log.info(
+                    "config.trainer.auto_determine_num_nodes ignored because no SLURM or LSF detected."
+                )
+
+        # Update the kwargs with the additional trainer kwargs
+        _update_kwargs(**cast(Any, config.trainer.additional_lightning_kwargs))
+        _update_kwargs(**config.trainer.lightning_kwargs)
+        _update_kwargs(**kwargs_ctor)
+
+        return kwargs
+
+    @override
+    def __init__(
+        self,
+        config: BaseConfig,
+        /,
+        **kwargs: Unpack[LightningTrainerKwargs],
+    ):
+        self._ll_config = config
+        kwargs = self._update_kwargs(config, kwargs)
+        log.critical(f"LightningTrainer.__init__ with {kwargs=}.")
+
+        super().__init__(**kwargs)
+
+        # Replace the signal connector with our own.
+        self._signal_connector = _SignalConnector(self)
+
+        # Print out the log dir, so that we can easily find it in the logs.
+        if log_dir := self.log_dir:
+            log_dir = str(Path(log_dir).resolve())
+        log.critical(f"LightningTrainer log directory: {self.log_dir}.")
+
+        # Checkpoint loading
+        if (
+            ckpt_loading := self._ll_config.trainer.checkpoint_loading
+        ) and ckpt_loading.path:
+            self.ckpt_path = ckpt_loading.path
+
+    @contextlib.contextmanager
+    def _actsave_context(self, model: LightningModule):
+        hparams = cast(BaseConfig, model.hparams)
+        if not (actsave_config := hparams.trainer.actsave):
+            yield
+            return
+
+        # Enter actsave context
+        with ActSave.enabled(actsave_config.resolve_save_dir(hparams)):
+            yield
+
+    @override
+    def _run(
+        self, model: LightningModule, ckpt_path: str | Path | None = None
+    ) -> _EVALUATE_OUTPUT | _PREDICT_OUTPUT | None:
+        """
+        Two things done here:
+            1. Lightning doesn't support gradient clipping with manual optimization.
+            We patch the `Trainer._run` method to throw if gradient clipping is enabled
+            and `model.automatic_optimization` is False.
+
+            2. We actually set up actsave here.
+        """
+
+        if not model.automatic_optimization and (
+            self.gradient_clip_val is not None
+            or self.gradient_clip_algorithm is not None
+        ):
+            raise ValueError(
+                "Automatic gradient clipping is not supported with manual optimization. "
+                f"Please set {model.__class__.__name__}.automatic_optimization to True "
+                "or disable automatic gradient clipping. "
+            )
+
+        with self._actsave_context(model):
+            return super()._run(model, ckpt_path)

nshtrainer/typecheck.py

@@ -0,0 +1,144 @@
+import os
+from collections.abc import Sequence
+from logging import getLogger
+from typing import Any
+
+import numpy as np
+import torch
+from jaxtyping import BFloat16 as BFloat16
+from jaxtyping import Bool as Bool
+from jaxtyping import Complex as Complex
+from jaxtyping import Complex64 as Complex64
+from jaxtyping import Complex128 as Complex128
+from jaxtyping import Float as Float
+from jaxtyping import Float16 as Float16
+from jaxtyping import Float32 as Float32
+from jaxtyping import Float64 as Float64
+from jaxtyping import Inexact as Inexact
+from jaxtyping import Int as Int
+from jaxtyping import Int4 as Int4
+from jaxtyping import Int8 as Int8
+from jaxtyping import Int16 as Int16
+from jaxtyping import Int32 as Int32
+from jaxtyping import Int64 as Int64
+from jaxtyping import Integer as Integer
+from jaxtyping import Key as Key
+from jaxtyping import Num as Num
+from jaxtyping import Real as Real
+from jaxtyping import Shaped as Shaped
+from jaxtyping import UInt as UInt
+from jaxtyping import UInt4 as UInt4
+from jaxtyping import UInt8 as UInt8
+from jaxtyping import UInt16 as UInt16
+from jaxtyping import UInt32 as UInt32
+from jaxtyping import UInt64 as UInt64
+from jaxtyping._storage import get_shape_memo, shape_str
+from torch import Tensor as Tensor
+from torch.nn.parameter import Parameter as Parameter
+from typing_extensions import TypeVar
+
+log = getLogger(__name__)
+
+DISABLE_ENV_KEY = "LL_DISABLE_TYPECHECKING"
+
+
+def typecheck_modules(modules: Sequence[str]):
+    """
+    Typecheck the given modules using `jaxtyping`.
+
+    Args:
+        modules: Modules to typecheck.
+    """
+    # If `DISABLE_ENV_KEY` is set and the environment variable is set, skip
+    #   typechecking.
+    if DISABLE_ENV_KEY is not None and bool(int(os.environ.get(DISABLE_ENV_KEY, "0"))):
+        log.critical(
+            f"Type checking is disabled due to the environment variable {DISABLE_ENV_KEY}."
+        )
+        return
+
+    # Install the jaxtyping import hook for this module.
+    from jaxtyping import install_import_hook
+
+    install_import_hook(modules, "beartype.beartype")
+
+    log.critical(f"Type checking the following modules: {modules}")
+
+
+def typecheck_this_module(additional_modules: Sequence[str] = ()):
+    """
+    Typecheck the calling module and any additional modules using `jaxtyping`.
+
+    Args:
+        additional_modules: Additional modules to typecheck.
+    """
+    # Get the calling module's name.
+    # Here, we can just use beartype's internal implementation behind
+    # `beartype_this_package`.
+    from beartype._util.func.utilfuncframe import get_frame, get_frame_package_name
+
+    # Get the calling module's name.
+    assert get_frame is not None, "get_frame is None"
+    frame = get_frame(1)
+    assert frame is not None, "frame is None"
+    calling_module_name = get_frame_package_name(frame)
+
+    # Typecheck the calling module + any additional modules.
+    typecheck_modules((calling_module_name, *additional_modules))
+
+
+def _make_error_str(input: Any, t: Any) -> str:
+    error_components: list[str] = []
+    error_components.append("Type checking error:")
+    if hasattr(t, "__instancecheck_str__"):
+        error_components.append(t.__instancecheck_str__(input))
+    if torch.is_tensor(input):
+        try:
+            from lovely_tensors import lovely
+
+            error_components.append(repr(lovely(input)))
+        except BaseException:
+            error_components.append(repr(input.shape))
+    error_components.append(shape_str(get_shape_memo()))
+
+    return "\n".join(error_components)
+
+
+T = TypeVar("T", torch.Tensor, np.ndarray, infer_variance=True)
+
+"""
+Patch to jaxtyping:
+
+In `jaxtyping._import_hook`, we add:
+def _has_isinstance_or_tassert(func_def):
+    for node in ast.walk(func_def):
+        if isinstance(node, ast.Call):
+            if isinstance(node.func, ast.Name) and node.func.id == "isinstance":
+                return True
+            elif isinstance(node.func, ast.Name) and node.func.id == "tassert":
+                return True
+    return False
+
+and we check this when adding the decorators.
+"""
+
+
+def tassert(t: Any, input: T | tuple[T, ...]):
+    """
+    Typecheck the input against the given type.
+
+    Args:
+        t: Type to check against.
+        input: Input to check.
+    """
+
+    # Ignore typechecking if the environment variable is set.
+    if DISABLE_ENV_KEY is not None and bool(int(os.environ.get(DISABLE_ENV_KEY, "0"))):
+        return
+
+    if isinstance(input, tuple):
+        for i in input:
+            assert isinstance(i, t), _make_error_str(i, t)
+        return
+    else:
+        assert isinstance(input, t), _make_error_str(input, t)