Add files using upload-large-folder tool

LTA_openwebtext_dualt/logs/lm1b_classic_dirichlet_every1k_infer_watch/infer_lta_lm1b_classic_dirichlet_len256_gbs512_4gpu_10k_save1k_20260523_step_0016000_t1p45.log

@@ -0,0 +1,20 @@
+[watch-classic-1k] 2026-05-23_16:05:04 infer runs/lta_lm1b_classic_dirichlet_len256_gbs512_4gpu_10k_save1k_20260523/step_0016000.pt -> docs/lta_samples/metrics_20260523/lm1b_classic_dirichlet_len256_every1k_normal_steps_state_t1p45_c1024_n256/lta_lm1b_classic_dirichlet_len256_gbs512_4gpu_10k_save1k_20260523/step_0016000
+[ckpt] runs/lta_lm1b_classic_dirichlet_len256_gbs512_4gpu_10k_save1k_20260523/step_0016000.pt step=16000
+[decode] steps128_c1024_t1p45 generated 16/256
+[decode] steps128_c1024_t1p45 generated 32/256
+[decode] steps128_c1024_t1p45 generated 48/256
+[decode] steps128_c1024_t1p45 generated 64/256
+[decode] steps128_c1024_t1p45 generated 80/256
+[decode] steps128_c1024_t1p45 generated 96/256
+[decode] steps128_c1024_t1p45 generated 112/256
+[decode] steps128_c1024_t1p45 generated 128/256
+[decode] steps128_c1024_t1p45 generated 144/256
+[decode] steps128_c1024_t1p45 generated 160/256
+[decode] steps128_c1024_t1p45 generated 176/256
+[decode] steps128_c1024_t1p45 generated 192/256
+[decode] steps128_c1024_t1p45 generated 208/256
+[decode] steps128_c1024_t1p45 generated 224/256
+[decode] steps128_c1024_t1p45 generated 240/256
+[decode] steps128_c1024_t1p45 generated 256/256
+[summary] {"name": "steps128_c1024_t1p45", "step": 16000, "decode_steps": 128, "concentration_max": 1024.0, "raw_genppl": 13.810265002566238, "stripped_genppl": 13.59726076657564, "sample_entropy": 1.4121913855304282, "distinct_1": 0.0078125, "distinct_2": 0.028247549019607843, "top_token_mass": 0.387176513671875, "raw_kept": 256, "stripped_kept": 256}
+[watch-classic-1k] 2026-05-23_16:07:59 done step_0016000

LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/packaging/_elffile.py

@@ -0,0 +1,108 @@
+"""
+ELF file parser.
+
+This provides a class ``ELFFile`` that parses an ELF executable in a similar
+interface to ``ZipFile``. Only the read interface is implemented.
+
+ELF header: https://refspecs.linuxfoundation.org/elf/gabi4+/ch4.eheader.html
+"""
+
+from __future__ import annotations
+
+import enum
+import os
+import struct
+from typing import IO
+
+
+class ELFInvalid(ValueError):
+    pass
+
+
+class EIClass(enum.IntEnum):
+    C32 = 1
+    C64 = 2
+
+
+class EIData(enum.IntEnum):
+    Lsb = 1
+    Msb = 2
+
+
+class EMachine(enum.IntEnum):
+    I386 = 3
+    S390 = 22
+    Arm = 40
+    X8664 = 62
+    AArc64 = 183
+
+
+class ELFFile:
+    """
+    Representation of an ELF executable.
+    """
+
+    def __init__(self, f: IO[bytes]) -> None:
+        self._f = f
+
+        try:
+            ident = self._read("16B")
+        except struct.error as e:
+            raise ELFInvalid("unable to parse identification") from e
+        magic = bytes(ident[:4])
+        if magic != b"\x7fELF":
+            raise ELFInvalid(f"invalid magic: {magic!r}")
+
+        self.capacity = ident[4]  # Format for program header (bitness).
+        self.encoding = ident[5]  # Data structure encoding (endianness).
+
+        try:
+            # e_fmt: Format for program header.
+            # p_fmt: Format for section header.
+            # p_idx: Indexes to find p_type, p_offset, and p_filesz.
+            e_fmt, self._p_fmt, self._p_idx = {
+                (1, 1): ("<HHIIIIIHHH", "<IIIIIIII", (0, 1, 4)),  # 32-bit LSB.
+                (1, 2): (">HHIIIIIHHH", ">IIIIIIII", (0, 1, 4)),  # 32-bit MSB.
+                (2, 1): ("<HHIQQQIHHH", "<IIQQQQQQ", (0, 2, 5)),  # 64-bit LSB.
+                (2, 2): (">HHIQQQIHHH", ">IIQQQQQQ", (0, 2, 5)),  # 64-bit MSB.
+            }[(self.capacity, self.encoding)]
+        except KeyError as e:
+            raise ELFInvalid(
+                f"unrecognized capacity ({self.capacity}) or encoding ({self.encoding})"
+            ) from e
+
+        try:
+            (
+                _,
+                self.machine,  # Architecture type.
+                _,
+                _,
+                self._e_phoff,  # Offset of program header.
+                _,
+                self.flags,  # Processor-specific flags.
+                _,
+                self._e_phentsize,  # Size of section.
+                self._e_phnum,  # Number of sections.
+            ) = self._read(e_fmt)
+        except struct.error as e:
+            raise ELFInvalid("unable to parse machine and section information") from e
+
+    def _read(self, fmt: str) -> tuple[int, ...]:
+        return struct.unpack(fmt, self._f.read(struct.calcsize(fmt)))
+
+    @property
+    def interpreter(self) -> str | None:
+        """
+        The path recorded in the ``PT_INTERP`` section header.
+        """
+        for index in range(self._e_phnum):
+            self._f.seek(self._e_phoff + self._e_phentsize * index)
+            try:
+                data = self._read(self._p_fmt)
+            except struct.error:
+                continue
+            if data[self._p_idx[0]] != 3:  # Not PT_INTERP.
+                continue
+            self._f.seek(data[self._p_idx[1]])
+            return os.fsdecode(self._f.read(data[self._p_idx[2]])).strip("\0")
+        return None

LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/packaging/_manylinux.py

@@ -0,0 +1,262 @@
+from __future__ import annotations
+
+import collections
+import contextlib
+import functools
+import os
+import re
+import sys
+import warnings
+from typing import Generator, Iterator, NamedTuple, Sequence
+
+from ._elffile import EIClass, EIData, ELFFile, EMachine
+
+EF_ARM_ABIMASK = 0xFF000000
+EF_ARM_ABI_VER5 = 0x05000000
+EF_ARM_ABI_FLOAT_HARD = 0x00000400
+
+_ALLOWED_ARCHS = {
+    "x86_64",
+    "aarch64",
+    "ppc64",
+    "ppc64le",
+    "s390x",
+    "loongarch64",
+    "riscv64",
+}
+
+
+# `os.PathLike` not a generic type until Python 3.9, so sticking with `str`
+# as the type for `path` until then.
+@contextlib.contextmanager
+def _parse_elf(path: str) -> Generator[ELFFile | None, None, None]:
+    try:
+        with open(path, "rb") as f:
+            yield ELFFile(f)
+    except (OSError, TypeError, ValueError):
+        yield None
+
+
+def _is_linux_armhf(executable: str) -> bool:
+    # hard-float ABI can be detected from the ELF header of the running
+    # process
+    # https://static.docs.arm.com/ihi0044/g/aaelf32.pdf
+    with _parse_elf(executable) as f:
+        return (
+            f is not None
+            and f.capacity == EIClass.C32
+            and f.encoding == EIData.Lsb
+            and f.machine == EMachine.Arm
+            and f.flags & EF_ARM_ABIMASK == EF_ARM_ABI_VER5
+            and f.flags & EF_ARM_ABI_FLOAT_HARD == EF_ARM_ABI_FLOAT_HARD
+        )
+
+
+def _is_linux_i686(executable: str) -> bool:
+    with _parse_elf(executable) as f:
+        return (
+            f is not None
+            and f.capacity == EIClass.C32
+            and f.encoding == EIData.Lsb
+            and f.machine == EMachine.I386
+        )
+
+
+def _have_compatible_abi(executable: str, archs: Sequence[str]) -> bool:
+    if "armv7l" in archs:
+        return _is_linux_armhf(executable)
+    if "i686" in archs:
+        return _is_linux_i686(executable)
+    return any(arch in _ALLOWED_ARCHS for arch in archs)
+
+
+# If glibc ever changes its major version, we need to know what the last
+# minor version was, so we can build the complete list of all versions.
+# For now, guess what the highest minor version might be, assume it will
+# be 50 for testing. Once this actually happens, update the dictionary
+# with the actual value.
+_LAST_GLIBC_MINOR: dict[int, int] = collections.defaultdict(lambda: 50)
+
+
+class _GLibCVersion(NamedTuple):
+    major: int
+    minor: int
+
+
+def _glibc_version_string_confstr() -> str | None:
+    """
+    Primary implementation of glibc_version_string using os.confstr.
+    """
+    # os.confstr is quite a bit faster than ctypes.DLL. It's also less likely
+    # to be broken or missing. This strategy is used in the standard library
+    # platform module.
+    # https://github.com/python/cpython/blob/fcf1d003bf4f0100c/Lib/platform.py#L175-L183
+    try:
+        # Should be a string like "glibc 2.17".
+        version_string: str | None = os.confstr("CS_GNU_LIBC_VERSION")
+        assert version_string is not None
+        _, version = version_string.rsplit()
+    except (AssertionError, AttributeError, OSError, ValueError):
+        # os.confstr() or CS_GNU_LIBC_VERSION not available (or a bad value)...
+        return None
+    return version
+
+
+def _glibc_version_string_ctypes() -> str | None:
+    """
+    Fallback implementation of glibc_version_string using ctypes.
+    """
+    try:
+        import ctypes  # noqa: PLC0415
+    except ImportError:
+        return None
+
+    # ctypes.CDLL(None) internally calls dlopen(NULL), and as the dlopen
+    # manpage says, "If filename is NULL, then the returned handle is for the
+    # main program". This way we can let the linker do the work to figure out
+    # which libc our process is actually using.
+    #
+    # We must also handle the special case where the executable is not a
+    # dynamically linked executable. This can occur when using musl libc,
+    # for example. In this situation, dlopen() will error, leading to an
+    # OSError. Interestingly, at least in the case of musl, there is no
+    # errno set on the OSError. The single string argument used to construct
+    # OSError comes from libc itself and is therefore not portable to
+    # hard code here. In any case, failure to call dlopen() means we
+    # can proceed, so we bail on our attempt.
+    try:
+        process_namespace = ctypes.CDLL(None)
+    except OSError:
+        return None
+
+    try:
+        gnu_get_libc_version = process_namespace.gnu_get_libc_version
+    except AttributeError:
+        # Symbol doesn't exist -> therefore, we are not linked to
+        # glibc.
+        return None
+
+    # Call gnu_get_libc_version, which returns a string like "2.5"
+    gnu_get_libc_version.restype = ctypes.c_char_p
+    version_str: str = gnu_get_libc_version()
+    # py2 / py3 compatibility:
+    if not isinstance(version_str, str):
+        version_str = version_str.decode("ascii")
+
+    return version_str
+
+
+def _glibc_version_string() -> str | None:
+    """Returns glibc version string, or None if not using glibc."""
+    return _glibc_version_string_confstr() or _glibc_version_string_ctypes()
+
+
+def _parse_glibc_version(version_str: str) -> _GLibCVersion:
+    """Parse glibc version.
+
+    We use a regexp instead of str.split because we want to discard any
+    random junk that might come after the minor version -- this might happen
+    in patched/forked versions of glibc (e.g. Linaro's version of glibc
+    uses version strings like "2.20-2014.11"). See gh-3588.
+    """
+    m = re.match(r"(?P<major>[0-9]+)\.(?P<minor>[0-9]+)", version_str)
+    if not m:
+        warnings.warn(
+            f"Expected glibc version with 2 components major.minor, got: {version_str}",
+            RuntimeWarning,
+            stacklevel=2,
+        )
+        return _GLibCVersion(-1, -1)
+    return _GLibCVersion(int(m.group("major")), int(m.group("minor")))
+
+
+@functools.lru_cache
+def _get_glibc_version() -> _GLibCVersion:
+    version_str = _glibc_version_string()
+    if version_str is None:
+        return _GLibCVersion(-1, -1)
+    return _parse_glibc_version(version_str)
+
+
+# From PEP 513, PEP 600
+def _is_compatible(arch: str, version: _GLibCVersion) -> bool:
+    sys_glibc = _get_glibc_version()
+    if sys_glibc < version:
+        return False
+    # Check for presence of _manylinux module.
+    try:
+        import _manylinux  # noqa: PLC0415
+    except ImportError:
+        return True
+    if hasattr(_manylinux, "manylinux_compatible"):
+        result = _manylinux.manylinux_compatible(version[0], version[1], arch)
+        if result is not None:
+            return bool(result)
+        return True
+    if version == _GLibCVersion(2, 5) and hasattr(_manylinux, "manylinux1_compatible"):
+        return bool(_manylinux.manylinux1_compatible)
+    if version == _GLibCVersion(2, 12) and hasattr(
+        _manylinux, "manylinux2010_compatible"
+    ):
+        return bool(_manylinux.manylinux2010_compatible)
+    if version == _GLibCVersion(2, 17) and hasattr(
+        _manylinux, "manylinux2014_compatible"
+    ):
+        return bool(_manylinux.manylinux2014_compatible)
+    return True
+
+
+_LEGACY_MANYLINUX_MAP: dict[_GLibCVersion, str] = {
+    # CentOS 7 w/ glibc 2.17 (PEP 599)
+    _GLibCVersion(2, 17): "manylinux2014",
+    # CentOS 6 w/ glibc 2.12 (PEP 571)
+    _GLibCVersion(2, 12): "manylinux2010",
+    # CentOS 5 w/ glibc 2.5 (PEP 513)
+    _GLibCVersion(2, 5): "manylinux1",
+}
+
+
+def platform_tags(archs: Sequence[str]) -> Iterator[str]:
+    """Generate manylinux tags compatible to the current platform.
+
+    :param archs: Sequence of compatible architectures.
+        The first one shall be the closest to the actual architecture and be the part of
+        platform tag after the ``linux_`` prefix, e.g. ``x86_64``.
+        The ``linux_`` prefix is assumed as a prerequisite for the current platform to
+        be manylinux-compatible.
+
+    :returns: An iterator of compatible manylinux tags.
+    """
+    if not _have_compatible_abi(sys.executable, archs):
+        return
+    # Oldest glibc to be supported regardless of architecture is (2, 17).
+    too_old_glibc2 = _GLibCVersion(2, 16)
+    if set(archs) & {"x86_64", "i686"}:
+        # On x86/i686 also oldest glibc to be supported is (2, 5).
+        too_old_glibc2 = _GLibCVersion(2, 4)
+    current_glibc = _GLibCVersion(*_get_glibc_version())
+    glibc_max_list = [current_glibc]
+    # We can assume compatibility across glibc major versions.
+    # https://sourceware.org/bugzilla/show_bug.cgi?id=24636
+    #
+    # Build a list of maximum glibc versions so that we can
+    # output the canonical list of all glibc from current_glibc
+    # down to too_old_glibc2, including all intermediary versions.
+    for glibc_major in range(current_glibc.major - 1, 1, -1):
+        glibc_minor = _LAST_GLIBC_MINOR[glibc_major]
+        glibc_max_list.append(_GLibCVersion(glibc_major, glibc_minor))
+    for arch in archs:
+        for glibc_max in glibc_max_list:
+            if glibc_max.major == too_old_glibc2.major:
+                min_minor = too_old_glibc2.minor
+            else:
+                # For other glibc major versions oldest supported is (x, 0).
+                min_minor = -1
+            for glibc_minor in range(glibc_max.minor, min_minor, -1):
+                glibc_version = _GLibCVersion(glibc_max.major, glibc_minor)
+                if _is_compatible(arch, glibc_version):
+                    yield "manylinux_{}_{}_{}".format(*glibc_version, arch)
+
+                    # Handle the legacy manylinux1, manylinux2010, manylinux2014 tags.
+                    if legacy_tag := _LEGACY_MANYLINUX_MAP.get(glibc_version):
+                        yield f"{legacy_tag}_{arch}"

LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/packaging/_musllinux.py

@@ -0,0 +1,85 @@
+"""PEP 656 support.
+
+This module implements logic to detect if the currently running Python is
+linked against musl, and what musl version is used.
+"""
+
+from __future__ import annotations
+
+import functools
+import re
+import subprocess
+import sys
+from typing import Iterator, NamedTuple, Sequence
+
+from ._elffile import ELFFile
+
+
+class _MuslVersion(NamedTuple):
+    major: int
+    minor: int
+
+
+def _parse_musl_version(output: str) -> _MuslVersion | None:
+    lines = [n for n in (n.strip() for n in output.splitlines()) if n]
+    if len(lines) < 2 or lines[0][:4] != "musl":
+        return None
+    m = re.match(r"Version (\d+)\.(\d+)", lines[1])
+    if not m:
+        return None
+    return _MuslVersion(major=int(m.group(1)), minor=int(m.group(2)))
+
+
+@functools.lru_cache
+def _get_musl_version(executable: str) -> _MuslVersion | None:
+    """Detect currently-running musl runtime version.
+
+    This is done by checking the specified executable's dynamic linking
+    information, and invoking the loader to parse its output for a version
+    string. If the loader is musl, the output would be something like::
+
+        musl libc (x86_64)
+        Version 1.2.2
+        Dynamic Program Loader
+    """
+    try:
+        with open(executable, "rb") as f:
+            ld = ELFFile(f).interpreter
+    except (OSError, TypeError, ValueError):
+        return None
+    if ld is None or "musl" not in ld:
+        return None
+    proc = subprocess.run([ld], check=False, stderr=subprocess.PIPE, text=True)
+    return _parse_musl_version(proc.stderr)
+
+
+def platform_tags(archs: Sequence[str]) -> Iterator[str]:
+    """Generate musllinux tags compatible to the current platform.
+
+    :param archs: Sequence of compatible architectures.
+        The first one shall be the closest to the actual architecture and be the part of
+        platform tag after the ``linux_`` prefix, e.g. ``x86_64``.
+        The ``linux_`` prefix is assumed as a prerequisite for the current platform to
+        be musllinux-compatible.
+
+    :returns: An iterator of compatible musllinux tags.
+    """
+    sys_musl = _get_musl_version(sys.executable)
+    if sys_musl is None:  # Python not dynamically linked against musl.
+        return
+    for arch in archs:
+        for minor in range(sys_musl.minor, -1, -1):
+            yield f"musllinux_{sys_musl.major}_{minor}_{arch}"
+
+
+if __name__ == "__main__":  # pragma: no cover
+    import sysconfig
+
+    plat = sysconfig.get_platform()
+    assert plat.startswith("linux-"), "not linux"
+
+    print("plat:", plat)
+    print("musl:", _get_musl_version(sys.executable))
+    print("tags:", end=" ")
+    for t in platform_tags(re.sub(r"[.-]", "_", plat.split("-", 1)[-1])):
+        print(t, end="\n      ")

LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/packaging/_parser.py

@@ -0,0 +1,393 @@
+"""Handwritten parser of dependency specifiers.
+
+The docstring for each __parse_* function contains EBNF-inspired grammar representing
+the implementation.
+"""
+
+from __future__ import annotations
+
+import ast
+from typing import List, Literal, NamedTuple, Sequence, Tuple, Union
+
+from ._tokenizer import DEFAULT_RULES, Tokenizer
+
+
+class Node:
+    __slots__ = ("value",)
+
+    def __init__(self, value: str) -> None:
+        self.value = value
+
+    def __str__(self) -> str:
+        return self.value
+
+    def __repr__(self) -> str:
+        return f"<{self.__class__.__name__}({self.value!r})>"
+
+    def serialize(self) -> str:
+        raise NotImplementedError
+
+    def __getstate__(self) -> str:
+        # Return just the value string for compactness and stability.
+        return self.value
+
+    def _restore_value(self, value: object) -> None:
+        if not isinstance(value, str):
+            raise TypeError(
+                f"Cannot restore {self.__class__.__name__} value from {value!r}"
+            )
+        self.value = value
+
+    def __setstate__(self, state: object) -> None:
+        if isinstance(state, str):
+            # New format (26.2+): just the value string.
+            self._restore_value(state)
+            return
+        if isinstance(state, tuple) and len(state) == 2:
+            # Old format (packaging <= 26.0, __slots__): (None, {slot: value}).
+            _, slot_dict = state
+            if isinstance(slot_dict, dict) and "value" in slot_dict:
+                self._restore_value(slot_dict["value"])
+                return
+        if isinstance(state, dict) and "value" in state:
+            # Old format (packaging <= 25.0, no __slots__): plain __dict__.
+            self._restore_value(state["value"])
+            return
+        raise TypeError(f"Cannot restore {self.__class__.__name__} from {state!r}")
+
+
+class Variable(Node):
+    __slots__ = ()
+
+    def serialize(self) -> str:
+        return str(self)
+
+
+class Value(Node):
+    __slots__ = ()
+
+    def serialize(self) -> str:
+        return f'"{self}"'
+
+
+class Op(Node):
+    __slots__ = ()
+
+    def serialize(self) -> str:
+        return str(self)
+
+
+MarkerLogical = Literal["and", "or"]
+MarkerVar = Union[Variable, Value]
+MarkerItem = Tuple[MarkerVar, Op, MarkerVar]
+MarkerAtom = Union[MarkerItem, Sequence["MarkerAtom"]]
+MarkerList = List[Union["MarkerList", MarkerAtom, MarkerLogical]]
+
+
+class ParsedRequirement(NamedTuple):
+    name: str
+    url: str
+    extras: list[str]
+    specifier: str
+    marker: MarkerList | None
+
+
+# --------------------------------------------------------------------------------------
+# Recursive descent parser for dependency specifier
+# --------------------------------------------------------------------------------------
+def parse_requirement(source: str) -> ParsedRequirement:
+    return _parse_requirement(Tokenizer(source, rules=DEFAULT_RULES))
+
+
+def _parse_requirement(tokenizer: Tokenizer) -> ParsedRequirement:
+    """
+    requirement = WS? IDENTIFIER WS? extras WS? requirement_details
+    """
+    tokenizer.consume("WS")
+
+    name_token = tokenizer.expect(
+        "IDENTIFIER", expected="package name at the start of dependency specifier"
+    )
+    name = name_token.text
+    tokenizer.consume("WS")
+
+    extras = _parse_extras(tokenizer)
+    tokenizer.consume("WS")
+
+    url, specifier, marker = _parse_requirement_details(tokenizer)
+    tokenizer.expect("END", expected="end of dependency specifier")
+
+    return ParsedRequirement(name, url, extras, specifier, marker)
+
+
+def _parse_requirement_details(
+    tokenizer: Tokenizer,
+) -> tuple[str, str, MarkerList | None]:
+    """
+    requirement_details = AT URL (WS requirement_marker?)?
+                        | specifier WS? (requirement_marker)?
+    """
+
+    specifier = ""
+    url = ""
+    marker = None
+
+    if tokenizer.check("AT"):
+        tokenizer.read()
+        tokenizer.consume("WS")
+
+        url_start = tokenizer.position
+        url = tokenizer.expect("URL", expected="URL after @").text
+        if tokenizer.check("END", peek=True):
+            return (url, specifier, marker)
+
+        tokenizer.expect("WS", expected="whitespace after URL")
+
+        # The input might end after whitespace.
+        if tokenizer.check("END", peek=True):
+            return (url, specifier, marker)
+
+        marker = _parse_requirement_marker(
+            tokenizer,
+            span_start=url_start,
+            expected="semicolon (after URL and whitespace)",
+        )
+    else:
+        specifier_start = tokenizer.position
+        specifier = _parse_specifier(tokenizer)
+        tokenizer.consume("WS")
+
+        if tokenizer.check("END", peek=True):
+            return (url, specifier, marker)
+
+        marker = _parse_requirement_marker(
+            tokenizer,
+            span_start=specifier_start,
+            expected=(
+                "comma (within version specifier), semicolon (after version specifier)"
+                if specifier
+                else "semicolon (after name with no version specifier)"
+            ),
+        )
+
+    return (url, specifier, marker)
+
+
+def _parse_requirement_marker(
+    tokenizer: Tokenizer, *, span_start: int, expected: str
+) -> MarkerList:
+    """
+    requirement_marker = SEMICOLON marker WS?
+    """
+
+    if not tokenizer.check("SEMICOLON"):
+        tokenizer.raise_syntax_error(
+            f"Expected {expected} or end",
+            span_start=span_start,
+            span_end=None,
+        )
+    tokenizer.read()
+
+    marker = _parse_marker(tokenizer)
+    tokenizer.consume("WS")
+
+    return marker
+
+
+def _parse_extras(tokenizer: Tokenizer) -> list[str]:
+    """
+    extras = (LEFT_BRACKET wsp* extras_list? wsp* RIGHT_BRACKET)?
+    """
+    if not tokenizer.check("LEFT_BRACKET", peek=True):
+        return []
+
+    with tokenizer.enclosing_tokens(
+        "LEFT_BRACKET",
+        "RIGHT_BRACKET",
+        around="extras",
+    ):
+        tokenizer.consume("WS")
+        extras = _parse_extras_list(tokenizer)
+        tokenizer.consume("WS")
+
+    return extras
+
+
+def _parse_extras_list(tokenizer: Tokenizer) -> list[str]:
+    """
+    extras_list = identifier (wsp* ',' wsp* identifier)*
+    """
+    extras: list[str] = []
+
+    if not tokenizer.check("IDENTIFIER"):
+        return extras
+
+    extras.append(tokenizer.read().text)
+
+    while True:
+        tokenizer.consume("WS")
+        if tokenizer.check("IDENTIFIER", peek=True):
+            tokenizer.raise_syntax_error("Expected comma between extra names")
+        elif not tokenizer.check("COMMA"):
+            break
+
+        tokenizer.read()
+        tokenizer.consume("WS")
+
+        extra_token = tokenizer.expect("IDENTIFIER", expected="extra name after comma")
+        extras.append(extra_token.text)
+
+    return extras
+
+
+def _parse_specifier(tokenizer: Tokenizer) -> str:
+    """
+    specifier = LEFT_PARENTHESIS WS? version_many WS? RIGHT_PARENTHESIS
+              | WS? version_many WS?
+    """
+    with tokenizer.enclosing_tokens(
+        "LEFT_PARENTHESIS",
+        "RIGHT_PARENTHESIS",
+        around="version specifier",
+    ):
+        tokenizer.consume("WS")
+        parsed_specifiers = _parse_version_many(tokenizer)
+        tokenizer.consume("WS")
+
+    return parsed_specifiers
+
+
+def _parse_version_many(tokenizer: Tokenizer) -> str:
+    """
+    version_many = (SPECIFIER (WS? COMMA WS? SPECIFIER)*)?
+    """
+    parsed_specifiers = ""
+    while tokenizer.check("SPECIFIER"):
+        span_start = tokenizer.position
+        parsed_specifiers += tokenizer.read().text
+        if tokenizer.check("VERSION_PREFIX_TRAIL", peek=True):
+            tokenizer.raise_syntax_error(
+                ".* suffix can only be used with `==` or `!=` operators",
+                span_start=span_start,
+                span_end=tokenizer.position + 1,
+            )
+        if tokenizer.check("VERSION_LOCAL_LABEL_TRAIL", peek=True):
+            tokenizer.raise_syntax_error(
+                "Local version label can only be used with `==` or `!=` operators",
+                span_start=span_start,
+                span_end=tokenizer.position,
+            )
+        tokenizer.consume("WS")
+        if not tokenizer.check("COMMA"):
+            break
+        parsed_specifiers += tokenizer.read().text
+        tokenizer.consume("WS")
+
+    return parsed_specifiers
+
+
+# --------------------------------------------------------------------------------------
+# Recursive descent parser for marker expression
+# --------------------------------------------------------------------------------------
+def parse_marker(source: str) -> MarkerList:
+    return _parse_full_marker(Tokenizer(source, rules=DEFAULT_RULES))
+
+
+def _parse_full_marker(tokenizer: Tokenizer) -> MarkerList:
+    retval = _parse_marker(tokenizer)
+    tokenizer.expect("END", expected="end of marker expression")
+    return retval
+
+
+def _parse_marker(tokenizer: Tokenizer) -> MarkerList:
+    """
+    marker = marker_atom (BOOLOP marker_atom)+
+    """
+    expression = [_parse_marker_atom(tokenizer)]
+    while tokenizer.check("BOOLOP"):
+        token = tokenizer.read()
+        expr_right = _parse_marker_atom(tokenizer)
+        expression.extend((token.text, expr_right))
+    return expression
+
+
+def _parse_marker_atom(tokenizer: Tokenizer) -> MarkerAtom:
+    """
+    marker_atom = WS? LEFT_PARENTHESIS WS? marker WS? RIGHT_PARENTHESIS WS?
+                | WS? marker_item WS?
+    """
+
+    tokenizer.consume("WS")
+    if tokenizer.check("LEFT_PARENTHESIS", peek=True):
+        with tokenizer.enclosing_tokens(
+            "LEFT_PARENTHESIS",
+            "RIGHT_PARENTHESIS",
+            around="marker expression",
+        ):
+            tokenizer.consume("WS")
+            marker: MarkerAtom = _parse_marker(tokenizer)
+            tokenizer.consume("WS")
+    else:
+        marker = _parse_marker_item(tokenizer)
+    tokenizer.consume("WS")
+    return marker
+
+
+def _parse_marker_item(tokenizer: Tokenizer) -> MarkerItem:
+    """
+    marker_item = WS? marker_var WS? marker_op WS? marker_var WS?
+    """
+    tokenizer.consume("WS")
+    marker_var_left = _parse_marker_var(tokenizer)
+    tokenizer.consume("WS")
+    marker_op = _parse_marker_op(tokenizer)
+    tokenizer.consume("WS")
+    marker_var_right = _parse_marker_var(tokenizer)
+    tokenizer.consume("WS")
+    return (marker_var_left, marker_op, marker_var_right)
+
+
+def _parse_marker_var(tokenizer: Tokenizer) -> MarkerVar:  # noqa: RET503
+    """
+    marker_var = VARIABLE | QUOTED_STRING
+    """
+    if tokenizer.check("VARIABLE"):
+        return process_env_var(tokenizer.read().text.replace(".", "_"))
+    elif tokenizer.check("QUOTED_STRING"):
+        return process_python_str(tokenizer.read().text)
+    else:
+        tokenizer.raise_syntax_error(
+            message="Expected a marker variable or quoted string"
+        )
+
+
+def process_env_var(env_var: str) -> Variable:
+    if env_var in ("platform_python_implementation", "python_implementation"):
+        return Variable("platform_python_implementation")
+    else:
+        return Variable(env_var)
+
+
+def process_python_str(python_str: str) -> Value:
+    value = ast.literal_eval(python_str)
+    return Value(str(value))
+
+
+def _parse_marker_op(tokenizer: Tokenizer) -> Op:
+    """
+    marker_op = IN | NOT IN | OP
+    """
+    if tokenizer.check("IN"):
+        tokenizer.read()
+        return Op("in")
+    elif tokenizer.check("NOT"):
+        tokenizer.read()
+        tokenizer.expect("WS", expected="whitespace after 'not'")
+        tokenizer.expect("IN", expected="'in' after 'not'")
+        return Op("not in")
+    elif tokenizer.check("OP"):
+        return Op(tokenizer.read().text)
+    else:
+        return tokenizer.raise_syntax_error(
+            "Expected marker operator, one of <=, <, !=, ==, >=, >, ~=, ===, in, not in"
+        )

LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/packaging/dependency_groups.py

@@ -0,0 +1,302 @@
+from __future__ import annotations
+
+import re
+from collections.abc import Mapping, Sequence
+
+from .errors import _ErrorCollector
+from .requirements import Requirement
+
+__all__ = [
+    "CyclicDependencyGroup",
+    "DependencyGroupInclude",
+    "DependencyGroupResolver",
+    "DuplicateGroupNames",
+    "InvalidDependencyGroupObject",
+    "resolve_dependency_groups",
+]
+
+
+def __dir__() -> list[str]:
+    return __all__
+
+
+# -----------
+# Error Types
+# -----------
+
+
+class DuplicateGroupNames(ValueError):
+    """
+    The same dependency groups were defined twice, with different non-normalized names.
+    """
+
+
+class CyclicDependencyGroup(ValueError):
+    """
+    The dependency group includes form a cycle.
+    """
+
+    def __init__(self, requested_group: str, group: str, include_group: str) -> None:
+        self.requested_group = requested_group
+        self.group = group
+        self.include_group = include_group
+
+        if include_group == group:
+            reason = f"{group} includes itself"
+        else:
+            reason = f"{include_group} -> {group}, {group} -> {include_group}"
+        super().__init__(
+            "Cyclic dependency group include while resolving "
+            f"{requested_group}: {reason}"
+        )
+
+
+# in the PEP 735 spec, the tables in dependency group lists were described as
+# "Dependency Object Specifiers", but the only defined type of object was a
+# "Dependency Group Include" -- hence the naming of this error as "Object"
+class InvalidDependencyGroupObject(ValueError):
+    """
+    A member of a dependency group was identified as a dict, but was not in a valid
+    format.
+    """
+
+
+# ------------------------
+# Object Model & Interface
+# ------------------------
+
+
+class DependencyGroupInclude:
+    __slots__ = ("include_group",)
+
+    def __init__(self, include_group: str) -> None:
+        """
+        Initialize a DependencyGroupInclude.
+
+        :param include_group: The name of the group referred to by this include.
+        """
+        self.include_group = include_group
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}({self.include_group!r})"
+
+
+class DependencyGroupResolver:
+    """
+    A resolver for Dependency Group data.
+
+    This class handles caching, name normalization, cycle detection, and other
+    parsing requirements. There are only two public methods for exploring the data:
+    ``lookup()`` and ``resolve()``.
+
+    :param dependency_groups: A mapping, as provided via pyproject
+        ``[dependency-groups]``.
+    """
+
+    def __init__(
+        self,
+        dependency_groups: Mapping[str, Sequence[str | Mapping[str, str]]],
+    ) -> None:
+        errors = _ErrorCollector()
+
+        self.dependency_groups = _normalize_group_names(dependency_groups, errors)
+
+        # a map of group names to parsed data
+        self._parsed_groups: dict[
+            str, tuple[Requirement | DependencyGroupInclude, ...]
+        ] = {}
+        # a map of group names to their ancestors, used for cycle detection
+        self._include_graph_ancestors: dict[str, tuple[str, ...]] = {}
+        # a cache of completed resolutions to Requirement lists
+        self._resolve_cache: dict[str, tuple[Requirement, ...]] = {}
+
+        errors.finalize("[dependency-groups] data was invalid")
+
+    def lookup(self, group: str) -> tuple[Requirement | DependencyGroupInclude, ...]:
+        """
+        Lookup a group name, returning the parsed dependency data for that group.
+        This will not resolve includes.
+
+        :param group: the name of the group to lookup
+        """
+        group = _normalize_name(group)
+
+        with _ErrorCollector().on_exit(
+            f"[dependency-groups] data for {group!r} was malformed"
+        ) as errors:
+            return self._parse_group(group, errors)
+
+    def resolve(self, group: str) -> tuple[Requirement, ...]:
+        """
+        Resolve a dependency group to a list of requirements.
+
+        :param group: the name of the group to resolve
+        """
+        group = _normalize_name(group)
+
+        with _ErrorCollector().on_exit(
+            f"[dependency-groups] data for {group!r} was malformed"
+        ) as errors:
+            return self._resolve(group, group, errors)
+
+    def _resolve(
+        self, group: str, requested_group: str, errors: _ErrorCollector
+    ) -> tuple[Requirement, ...]:
+        """
+        This is a helper for cached resolution to strings. It preserves the name of the
+        group which the user initially requested in order to present a clearer error in
+        the event that a cycle is detected.
+
+        :param group: The normalized name of the group to resolve.
+        :param requested_group: The group which was used in the original, user-facing
+            request.
+        """
+        if group in self._resolve_cache:
+            return self._resolve_cache[group]
+
+        parsed = self._parse_group(group, errors)
+
+        resolved_group = []
+
+        for item in parsed:
+            if isinstance(item, Requirement):
+                resolved_group.append(item)
+            elif isinstance(item, DependencyGroupInclude):
+                include_group = _normalize_name(item.include_group)
+
+                # if a group is cyclic, record the error
+                # otherwise, follow the include_group reference
+                #
+                # this allows us to examine all includes in a group, even in the
+                # presence of errors
+                if include_group in self._include_graph_ancestors.get(group, ()):
+                    errors.error(
+                        CyclicDependencyGroup(
+                            requested_group, group, item.include_group
+                        )
+                    )
+                else:
+                    self._include_graph_ancestors[include_group] = (
+                        *self._include_graph_ancestors.get(group, ()),
+                        group,
+                    )
+                    resolved_group.extend(
+                        self._resolve(include_group, requested_group, errors)
+                    )
+            else:  # pragma: no cover
+                raise NotImplementedError(
+                    f"Invalid dependency group item after parse: {item}"
+                )
+
+        # in the event that errors were detected, present the group as empty and do not
+        # cache the result
+        # this ensures that repeated access to a cyclic group will raise multiple errors
+        if errors.errors:
+            return ()
+
+        self._resolve_cache[group] = tuple(resolved_group)
+        return self._resolve_cache[group]
+
+    def _parse_group(
+        self, group: str, errors: _ErrorCollector
+    ) -> tuple[Requirement | DependencyGroupInclude, ...]:
+        # short circuit -- never do the work twice
+        if group in self._parsed_groups:
+            return self._parsed_groups[group]
+
+        if group not in self.dependency_groups:
+            errors.error(LookupError(f"Dependency group '{group}' not found"))
+            return ()
+
+        raw_group = self.dependency_groups[group]
+        if isinstance(raw_group, str):
+            errors.error(
+                TypeError(
+                    f"Dependency group {group!r} contained a string rather than a list."
+                )
+            )
+            return ()
+
+        if not isinstance(raw_group, Sequence):
+            errors.error(
+                TypeError(f"Dependency group {group!r} is not a sequence type.")
+            )
+            return ()
+
+        elements: list[Requirement | DependencyGroupInclude] = []
+        for item in raw_group:
+            if isinstance(item, str):
+                # packaging.requirements.Requirement parsing ensures that this is a
+                # valid PEP 508 Dependency Specifier
+                # raises InvalidRequirement on failure
+                elements.append(Requirement(item))
+            elif isinstance(item, Mapping):
+                if tuple(item.keys()) != ("include-group",):
+                    errors.error(
+                        InvalidDependencyGroupObject(
+                            f"Invalid dependency group item: {item!r}"
+                        )
+                    )
+                else:
+                    include_group = item["include-group"]
+                    elements.append(DependencyGroupInclude(include_group=include_group))
+            else:
+                errors.error(TypeError(f"Invalid dependency group item: {item!r}"))
+
+        self._parsed_groups[group] = tuple(elements)
+        return self._parsed_groups[group]
+
+
+# --------------------
+# Functional Interface
+# --------------------
+
+
+def resolve_dependency_groups(
+    dependency_groups: Mapping[str, Sequence[str | Mapping[str, str]]], /, *groups: str
+) -> tuple[str, ...]:
+    """
+    Resolve a dependency group to a tuple of requirements, as strings.
+
+    :param dependency_groups: the parsed contents of the ``[dependency-groups]`` table
+        from ``pyproject.toml``
+    :param groups: the name of the group(s) to resolve
+    """
+    resolver = DependencyGroupResolver(dependency_groups)
+    return tuple(str(r) for group in groups for r in resolver.resolve(group))
+
+
+# ----------------
+# internal helpers
+# ----------------
+
+
+_NORMALIZE_PATTERN = re.compile(r"[-_.]+")
+
+
+def _normalize_name(name: str) -> str:
+    return _NORMALIZE_PATTERN.sub("-", name).lower()
+
+
+def _normalize_group_names(
+    dependency_groups: Mapping[str, Sequence[str | Mapping[str, str]]],
+    errors: _ErrorCollector,
+) -> dict[str, Sequence[str | Mapping[str, str]]]:
+    original_names: dict[str, list[str]] = {}
+    normalized_groups: dict[str, Sequence[str | Mapping[str, str]]] = {}
+
+    for group_name, value in dependency_groups.items():
+        normed_group_name = _normalize_name(group_name)
+        original_names.setdefault(normed_group_name, []).append(group_name)
+        normalized_groups[normed_group_name] = value
+
+    for normed_name, names in original_names.items():
+        if len(names) > 1:
+            errors.error(
+                DuplicateGroupNames(
+                    "Duplicate dependency group names: "
+                    f"{normed_name} ({', '.join(names)})"
+                )
+            )
+
+    return normalized_groups

LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/packaging/pylock.py

@@ -0,0 +1,905 @@
+from __future__ import annotations
+
+import dataclasses
+import logging
+import re
+from collections.abc import Mapping, Sequence
+from dataclasses import dataclass
+from datetime import datetime
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Protocol,
+    TypeVar,
+    cast,
+)
+from urllib.parse import urlparse
+
+from .markers import Environment, Marker, default_environment
+from .specifiers import SpecifierSet
+from .tags import create_compatible_tags_selector, sys_tags
+from .utils import (
+    NormalizedName,
+    is_normalized_name,
+    parse_sdist_filename,
+    parse_wheel_filename,
+)
+from .version import Version
+
+if TYPE_CHECKING:  # pragma: no cover
+    from collections.abc import Collection, Iterator
+    from pathlib import Path
+
+    from typing_extensions import Self
+
+    from .tags import Tag
+
+_logger = logging.getLogger(__name__)
+
+__all__ = [
+    "Package",
+    "PackageArchive",
+    "PackageDirectory",
+    "PackageSdist",
+    "PackageVcs",
+    "PackageWheel",
+    "Pylock",
+    "PylockUnsupportedVersionError",
+    "PylockValidationError",
+    "is_valid_pylock_path",
+]
+
+
+def __dir__() -> list[str]:
+    return __all__
+
+
+_T = TypeVar("_T")
+_T2 = TypeVar("_T2")
+
+
+class _FromMappingProtocol(Protocol):  # pragma: no cover
+    @classmethod
+    def _from_dict(cls, d: Mapping[str, Any]) -> Self: ...
+
+
+_FromMappingProtocolT = TypeVar("_FromMappingProtocolT", bound=_FromMappingProtocol)
+
+
+_PYLOCK_FILE_NAME_RE = re.compile(r"^pylock\.([^.]+)\.toml$")
+
+
+def is_valid_pylock_path(path: Path) -> bool:
+    """Check if the given path is a valid pylock file path."""
+    return path.name == "pylock.toml" or bool(_PYLOCK_FILE_NAME_RE.match(path.name))
+
+
+def _toml_key(key: str) -> str:
+    return key.replace("_", "-")
+
+
+def _toml_value(key: str, value: Any) -> Any:  # noqa: ANN401
+    if isinstance(value, (Version, Marker, SpecifierSet)):
+        return str(value)
+    if isinstance(value, Sequence) and key == "environments":
+        return [str(v) for v in value]
+    return value
+
+
+def _toml_dict_factory(data: list[tuple[str, Any]]) -> dict[str, Any]:
+    return {
+        _toml_key(key): _toml_value(key, value)
+        for key, value in data
+        if value is not None
+    }
+
+
+def _get(d: Mapping[str, Any], expected_type: type[_T], key: str) -> _T | None:
+    """Get a value from the dictionary and verify it's the expected type."""
+    if (value := d.get(key)) is None:
+        return None
+    if not isinstance(value, expected_type):
+        raise PylockValidationError(
+            f"Unexpected type {type(value).__name__} "
+            f"(expected {expected_type.__name__})",
+            context=key,
+        )
+    return value
+
+
+def _get_required(d: Mapping[str, Any], expected_type: type[_T], key: str) -> _T:
+    """Get a required value from the dictionary and verify it's the expected type."""
+    if (value := _get(d, expected_type, key)) is None:
+        raise _PylockRequiredKeyError(key)
+    return value
+
+
+def _get_sequence(
+    d: Mapping[str, Any], expected_item_type: type[_T], key: str
+) -> Sequence[_T] | None:
+    """Get a list value from the dictionary and verify it's the expected items type."""
+    if (value := _get(d, Sequence, key)) is None:  # type: ignore[type-abstract]
+        return None
+    if isinstance(value, (str, bytes)):
+        # special case: str and bytes are Sequences, but we want to reject it
+        raise PylockValidationError(
+            f"Unexpected type {type(value).__name__} (expected Sequence)",
+            context=key,
+        )
+    for i, item in enumerate(value):
+        if not isinstance(item, expected_item_type):
+            raise PylockValidationError(
+                f"Unexpected type {type(item).__name__} "
+                f"(expected {expected_item_type.__name__})",
+                context=f"{key}[{i}]",
+            )
+    return value
+
+
+def _get_as(
+    d: Mapping[str, Any],
+    expected_type: type[_T],
+    target_type: Callable[[_T], _T2],
+    key: str,
+) -> _T2 | None:
+    """Get a value from the dictionary, verify it's the expected type,
+    and convert to the target type.
+
+    This assumes the target_type constructor accepts the value.
+    """
+    if (value := _get(d, expected_type, key)) is None:
+        return None
+    try:
+        return target_type(value)
+    except Exception as e:
+        raise PylockValidationError(e, context=key) from e
+
+
+def _get_required_as(
+    d: Mapping[str, Any],
+    expected_type: type[_T],
+    target_type: Callable[[_T], _T2],
+    key: str,
+) -> _T2:
+    """Get a required value from the dict, verify it's the expected type,
+    and convert to the target type."""
+    if (value := _get_as(d, expected_type, target_type, key)) is None:
+        raise _PylockRequiredKeyError(key)
+    return value
+
+
+def _get_sequence_as(
+    d: Mapping[str, Any],
+    expected_item_type: type[_T],
+    target_item_type: Callable[[_T], _T2],
+    key: str,
+) -> list[_T2] | None:
+    """Get list value from dictionary and verify expected items type."""
+    if (value := _get_sequence(d, expected_item_type, key)) is None:
+        return None
+    result = []
+    try:
+        for item in value:
+            typed_item = target_item_type(item)
+            result.append(typed_item)
+    except Exception as e:
+        raise PylockValidationError(e, context=f"{key}[{len(result)}]") from e
+    return result
+
+
+def _get_object(
+    d: Mapping[str, Any], target_type: type[_FromMappingProtocolT], key: str
+) -> _FromMappingProtocolT | None:
+    """Get a dictionary value from the dictionary and convert it to a dataclass."""
+    if (value := _get(d, Mapping, key)) is None:  # type: ignore[type-abstract]
+        return None
+    try:
+        return target_type._from_dict(value)
+    except Exception as e:
+        raise PylockValidationError(e, context=key) from e
+
+
+def _get_sequence_of_objects(
+    d: Mapping[str, Any], target_item_type: type[_FromMappingProtocolT], key: str
+) -> list[_FromMappingProtocolT] | None:
+    """Get a list value from the dictionary and convert its items to a dataclass."""
+    if (value := _get_sequence(d, Mapping, key)) is None:  # type: ignore[type-abstract]
+        return None
+    result: list[_FromMappingProtocolT] = []
+    try:
+        for item in value:
+            typed_item = target_item_type._from_dict(item)
+            result.append(typed_item)
+    except Exception as e:
+        raise PylockValidationError(e, context=f"{key}[{len(result)}]") from e
+    return result
+
+
+def _get_required_sequence_of_objects(
+    d: Mapping[str, Any], target_item_type: type[_FromMappingProtocolT], key: str
+) -> Sequence[_FromMappingProtocolT]:
+    """Get a required list value from the dictionary and convert its items to a
+    dataclass."""
+    if (result := _get_sequence_of_objects(d, target_item_type, key)) is None:
+        raise _PylockRequiredKeyError(key)
+    return result
+
+
+def _validate_normalized_name(name: str) -> NormalizedName:
+    """Validate that a string is a NormalizedName."""
+    if not is_normalized_name(name):
+        raise PylockValidationError(f"Name {name!r} is not normalized")
+    return NormalizedName(name)
+
+
+def _validate_path_url(path: str | None, url: str | None) -> None:
+    if not path and not url:
+        raise PylockValidationError("path or url must be provided")
+
+
+def _path_name(path: str | None) -> str | None:
+    if not path:
+        return None
+    # If the path is relative it MAY use POSIX-style path separators explicitly
+    # for portability
+    if "/" in path:
+        return path.rsplit("/", 1)[-1]
+    elif "\\" in path:
+        return path.rsplit("\\", 1)[-1]
+    else:
+        return path
+
+
+def _url_name(url: str | None) -> str | None:
+    if not url:
+        return None
+    url_path = urlparse(url).path
+    return url_path.rsplit("/", 1)[-1]
+
+
+def _validate_hashes(hashes: Mapping[str, Any]) -> Mapping[str, Any]:
+    if not hashes:
+        raise PylockValidationError("At least one hash must be provided")
+    if not all(isinstance(hash_val, str) for hash_val in hashes.values()):
+        raise PylockValidationError("Hash values must be strings")
+    return hashes
+
+
+class PylockValidationError(Exception):
+    """Raised when when input data is not spec-compliant."""
+
+    context: str | None = None
+    message: str
+
+    def __init__(
+        self,
+        cause: str | Exception,
+        *,
+        context: str | None = None,
+    ) -> None:
+        if isinstance(cause, PylockValidationError):
+            if cause.context:
+                self.context = (
+                    f"{context}.{cause.context}" if context else cause.context
+                )
+            else:
+                self.context = context
+            self.message = cause.message
+        else:
+            self.context = context
+            self.message = str(cause)
+
+    def __str__(self) -> str:
+        if self.context:
+            return f"{self.message} in {self.context!r}"
+        return self.message
+
+
+class _PylockRequiredKeyError(PylockValidationError):
+    def __init__(self, key: str) -> None:
+        super().__init__("Missing required value", context=key)
+
+
+class PylockUnsupportedVersionError(PylockValidationError):
+    """Raised when encountering an unsupported `lock_version`."""
+
+
+class PylockSelectError(Exception):
+    """Base exception for errors raised by :meth:`Pylock.select`."""
+
+
+@dataclass(frozen=True, init=False)
+class PackageVcs:
+    type: str
+    url: str | None = None
+    path: str | None = None
+    requested_revision: str | None = None
+    commit_id: str  # type: ignore[misc]
+    subdirectory: str | None = None
+
+    def __init__(
+        self,
+        *,
+        type: str,
+        url: str | None = None,
+        path: str | None = None,
+        requested_revision: str | None = None,
+        commit_id: str,
+        subdirectory: str | None = None,
+    ) -> None:
+        # In Python 3.10+ make dataclass kw_only=True and remove __init__
+        object.__setattr__(self, "type", type)
+        object.__setattr__(self, "url", url)
+        object.__setattr__(self, "path", path)
+        object.__setattr__(self, "requested_revision", requested_revision)
+        object.__setattr__(self, "commit_id", commit_id)
+        object.__setattr__(self, "subdirectory", subdirectory)
+
+    @classmethod
+    def _from_dict(cls, d: Mapping[str, Any]) -> Self:
+        package_vcs = cls(
+            type=_get_required(d, str, "type"),
+            url=_get(d, str, "url"),
+            path=_get(d, str, "path"),
+            requested_revision=_get(d, str, "requested-revision"),
+            commit_id=_get_required(d, str, "commit-id"),
+            subdirectory=_get(d, str, "subdirectory"),
+        )
+        _validate_path_url(package_vcs.path, package_vcs.url)
+        return package_vcs
+
+
+@dataclass(frozen=True, init=False)
+class PackageDirectory:
+    path: str
+    editable: bool | None = None
+    subdirectory: str | None = None
+
+    def __init__(
+        self,
+        *,
+        path: str,
+        editable: bool | None = None,
+        subdirectory: str | None = None,
+    ) -> None:
+        # In Python 3.10+ make dataclass kw_only=True and remove __init__
+        object.__setattr__(self, "path", path)
+        object.__setattr__(self, "editable", editable)
+        object.__setattr__(self, "subdirectory", subdirectory)
+
+    @classmethod
+    def _from_dict(cls, d: Mapping[str, Any]) -> Self:
+        return cls(
+            path=_get_required(d, str, "path"),
+            editable=_get(d, bool, "editable"),
+            subdirectory=_get(d, str, "subdirectory"),
+        )
+
+
+@dataclass(frozen=True, init=False)
+class PackageArchive:
+    url: str | None = None
+    path: str | None = None
+    size: int | None = None
+    upload_time: datetime | None = None
+    hashes: Mapping[str, str]  # type: ignore[misc]
+    subdirectory: str | None = None
+
+    def __init__(
+        self,
+        *,
+        url: str | None = None,
+        path: str | None = None,
+        size: int | None = None,
+        upload_time: datetime | None = None,
+        hashes: Mapping[str, str],
+        subdirectory: str | None = None,
+    ) -> None:
+        # In Python 3.10+ make dataclass kw_only=True and remove __init__
+        object.__setattr__(self, "url", url)
+        object.__setattr__(self, "path", path)
+        object.__setattr__(self, "size", size)
+        object.__setattr__(self, "upload_time", upload_time)
+        object.__setattr__(self, "hashes", hashes)
+        object.__setattr__(self, "subdirectory", subdirectory)
+
+    @classmethod
+    def _from_dict(cls, d: Mapping[str, Any]) -> Self:
+        package_archive = cls(
+            url=_get(d, str, "url"),
+            path=_get(d, str, "path"),
+            size=_get(d, int, "size"),
+            upload_time=_get(d, datetime, "upload-time"),
+            hashes=_get_required_as(d, Mapping, _validate_hashes, "hashes"),  # type: ignore[type-abstract]
+            subdirectory=_get(d, str, "subdirectory"),
+        )
+        _validate_path_url(package_archive.path, package_archive.url)
+        return package_archive
+
+
+@dataclass(frozen=True, init=False)
+class PackageSdist:
+    name: str | None = None
+    upload_time: datetime | None = None
+    url: str | None = None
+    path: str | None = None
+    size: int | None = None
+    hashes: Mapping[str, str]  # type: ignore[misc]
+
+    def __init__(
+        self,
+        *,
+        name: str | None = None,
+        upload_time: datetime | None = None,
+        url: str | None = None,
+        path: str | None = None,
+        size: int | None = None,
+        hashes: Mapping[str, str],
+    ) -> None:
+        # In Python 3.10+ make dataclass kw_only=True and remove __init__
+        object.__setattr__(self, "name", name)
+        object.__setattr__(self, "upload_time", upload_time)
+        object.__setattr__(self, "url", url)
+        object.__setattr__(self, "path", path)
+        object.__setattr__(self, "size", size)
+        object.__setattr__(self, "hashes", hashes)
+
+    @classmethod
+    def _from_dict(cls, d: Mapping[str, Any]) -> Self:
+        package_sdist = cls(
+            name=_get(d, str, "name"),
+            upload_time=_get(d, datetime, "upload-time"),
+            url=_get(d, str, "url"),
+            path=_get(d, str, "path"),
+            size=_get(d, int, "size"),
+            hashes=_get_required_as(d, Mapping, _validate_hashes, "hashes"),  # type: ignore[type-abstract]
+        )
+        _validate_path_url(package_sdist.path, package_sdist.url)
+        return package_sdist
+
+    @property
+    def filename(self) -> str:
+        """Get the filename of the sdist."""
+        filename = self.name or _path_name(self.path) or _url_name(self.url)
+        if not filename:
+            raise PylockValidationError("Cannot determine sdist filename")
+        return filename
+
+
+@dataclass(frozen=True, init=False)
+class PackageWheel:
+    name: str | None = None
+    upload_time: datetime | None = None
+    url: str | None = None
+    path: str | None = None
+    size: int | None = None
+    hashes: Mapping[str, str]  # type: ignore[misc]
+
+    def __init__(
+        self,
+        *,
+        name: str | None = None,
+        upload_time: datetime | None = None,
+        url: str | None = None,
+        path: str | None = None,
+        size: int | None = None,
+        hashes: Mapping[str, str],
+    ) -> None:
+        # In Python 3.10+ make dataclass kw_only=True and remove __init__
+        object.__setattr__(self, "name", name)
+        object.__setattr__(self, "upload_time", upload_time)
+        object.__setattr__(self, "url", url)
+        object.__setattr__(self, "path", path)
+        object.__setattr__(self, "size", size)
+        object.__setattr__(self, "hashes", hashes)
+
+    @classmethod
+    def _from_dict(cls, d: Mapping[str, Any]) -> Self:
+        package_wheel = cls(
+            name=_get(d, str, "name"),
+            upload_time=_get(d, datetime, "upload-time"),
+            url=_get(d, str, "url"),
+            path=_get(d, str, "path"),
+            size=_get(d, int, "size"),
+            hashes=_get_required_as(d, Mapping, _validate_hashes, "hashes"),  # type: ignore[type-abstract]
+        )
+        _validate_path_url(package_wheel.path, package_wheel.url)
+        return package_wheel
+
+    @property
+    def filename(self) -> str:
+        """Get the filename of the wheel."""
+        filename = self.name or _path_name(self.path) or _url_name(self.url)
+        if not filename:
+            raise PylockValidationError("Cannot determine wheel filename")
+        return filename
+
+
+@dataclass(frozen=True, init=False)
+class Package:
+    name: NormalizedName
+    version: Version | None = None
+    marker: Marker | None = None
+    requires_python: SpecifierSet | None = None
+    dependencies: Sequence[Mapping[str, Any]] | None = None
+    vcs: PackageVcs | None = None
+    directory: PackageDirectory | None = None
+    archive: PackageArchive | None = None
+    index: str | None = None
+    sdist: PackageSdist | None = None
+    wheels: Sequence[PackageWheel] | None = None
+    attestation_identities: Sequence[Mapping[str, Any]] | None = None
+    tool: Mapping[str, Any] | None = None
+
+    def __init__(
+        self,
+        *,
+        name: NormalizedName,
+        version: Version | None = None,
+        marker: Marker | None = None,
+        requires_python: SpecifierSet | None = None,
+        dependencies: Sequence[Mapping[str, Any]] | None = None,
+        vcs: PackageVcs | None = None,
+        directory: PackageDirectory | None = None,
+        archive: PackageArchive | None = None,
+        index: str | None = None,
+        sdist: PackageSdist | None = None,
+        wheels: Sequence[PackageWheel] | None = None,
+        attestation_identities: Sequence[Mapping[str, Any]] | None = None,
+        tool: Mapping[str, Any] | None = None,
+    ) -> None:
+        # In Python 3.10+ make dataclass kw_only=True and remove __init__
+        object.__setattr__(self, "name", name)
+        object.__setattr__(self, "version", version)
+        object.__setattr__(self, "marker", marker)
+        object.__setattr__(self, "requires_python", requires_python)
+        object.__setattr__(self, "dependencies", dependencies)
+        object.__setattr__(self, "vcs", vcs)
+        object.__setattr__(self, "directory", directory)
+        object.__setattr__(self, "archive", archive)
+        object.__setattr__(self, "index", index)
+        object.__setattr__(self, "sdist", sdist)
+        object.__setattr__(self, "wheels", wheels)
+        object.__setattr__(self, "attestation_identities", attestation_identities)
+        object.__setattr__(self, "tool", tool)
+
+    @classmethod
+    def _from_dict(cls, d: Mapping[str, Any]) -> Self:
+        package = cls(
+            name=_get_required_as(d, str, _validate_normalized_name, "name"),
+            version=_get_as(d, str, Version, "version"),
+            requires_python=_get_as(d, str, SpecifierSet, "requires-python"),
+            dependencies=_get_sequence(d, Mapping, "dependencies"),  # type: ignore[type-abstract]
+            marker=_get_as(d, str, Marker, "marker"),
+            vcs=_get_object(d, PackageVcs, "vcs"),
+            directory=_get_object(d, PackageDirectory, "directory"),
+            archive=_get_object(d, PackageArchive, "archive"),
+            index=_get(d, str, "index"),
+            sdist=_get_object(d, PackageSdist, "sdist"),
+            wheels=_get_sequence_of_objects(d, PackageWheel, "wheels"),
+            attestation_identities=_get_sequence(d, Mapping, "attestation-identities"),  # type: ignore[type-abstract]
+            tool=_get(d, Mapping, "tool"),  # type: ignore[type-abstract]
+        )
+        distributions = bool(package.sdist) + len(package.wheels or [])
+        direct_urls = (
+            bool(package.vcs) + bool(package.directory) + bool(package.archive)
+        )
+        if distributions > 0 and direct_urls > 0:
+            raise PylockValidationError(
+                "None of vcs, directory, archive must be set if sdist or wheels are set"
+            )
+        if distributions == 0 and direct_urls != 1:
+            raise PylockValidationError(
+                "Exactly one of vcs, directory, archive must be set "
+                "if sdist and wheels are not set"
+            )
+        for i, wheel in enumerate(package.wheels or []):
+            try:
+                (name, version, _, _) = parse_wheel_filename(wheel.filename)
+            except Exception as e:
+                raise PylockValidationError(
+                    f"Invalid wheel filename {wheel.filename!r}",
+                    context=f"wheels[{i}]",
+                ) from e
+            if name != package.name:
+                raise PylockValidationError(
+                    f"Name in {wheel.filename!r} is not consistent with "
+                    f"package name {package.name!r}",
+                    context=f"wheels[{i}]",
+                )
+            if package.version and version != package.version:
+                raise PylockValidationError(
+                    f"Version in {wheel.filename!r} is not consistent with "
+                    f"package version {str(package.version)!r}",
+                    context=f"wheels[{i}]",
+                )
+        if package.sdist:
+            try:
+                name, version = parse_sdist_filename(package.sdist.filename)
+            except Exception as e:
+                raise PylockValidationError(
+                    f"Invalid sdist filename {package.sdist.filename!r}",
+                    context="sdist",
+                ) from e
+            if name != package.name:
+                raise PylockValidationError(
+                    f"Name in {package.sdist.filename!r} is not consistent with "
+                    f"package name {package.name!r}",
+                    context="sdist",
+                )
+            if package.version and version != package.version:
+                raise PylockValidationError(
+                    f"Version in {package.sdist.filename!r} is not consistent with "
+                    f"package version {str(package.version)!r}",
+                    context="sdist",
+                )
+        try:
+            for i, attestation_identity in enumerate(  # noqa: B007
+                package.attestation_identities or []
+            ):
+                _get_required(attestation_identity, str, "kind")
+        except Exception as e:
+            raise PylockValidationError(
+                e, context=f"attestation-identities[{i}]"
+            ) from e
+        return package
+
+    @property
+    def is_direct(self) -> bool:
+        return not (self.sdist or self.wheels)
+
+
+@dataclass(frozen=True, init=False)
+class Pylock:
+    """A class representing a pylock file."""
+
+    lock_version: Version
+    environments: Sequence[Marker] | None = None
+    requires_python: SpecifierSet | None = None
+    extras: Sequence[NormalizedName] | None = None
+    dependency_groups: Sequence[str] | None = None
+    default_groups: Sequence[str] | None = None
+    created_by: str  # type: ignore[misc]
+    packages: Sequence[Package]  # type: ignore[misc]
+    tool: Mapping[str, Any] | None = None
+
+    def __init__(
+        self,
+        *,
+        lock_version: Version,
+        environments: Sequence[Marker] | None = None,
+        requires_python: SpecifierSet | None = None,
+        extras: Sequence[NormalizedName] | None = None,
+        dependency_groups: Sequence[str] | None = None,
+        default_groups: Sequence[str] | None = None,
+        created_by: str,
+        packages: Sequence[Package],
+        tool: Mapping[str, Any] | None = None,
+    ) -> None:
+        # In Python 3.10+ make dataclass kw_only=True and remove __init__
+        object.__setattr__(self, "lock_version", lock_version)
+        object.__setattr__(self, "environments", environments)
+        object.__setattr__(self, "requires_python", requires_python)
+        object.__setattr__(self, "extras", extras)
+        object.__setattr__(self, "dependency_groups", dependency_groups)
+        object.__setattr__(self, "default_groups", default_groups)
+        object.__setattr__(self, "created_by", created_by)
+        object.__setattr__(self, "packages", packages)
+        object.__setattr__(self, "tool", tool)
+
+    @classmethod
+    def _from_dict(cls, d: Mapping[str, Any]) -> Self:
+        pylock = cls(
+            lock_version=_get_required_as(d, str, Version, "lock-version"),
+            environments=_get_sequence_as(d, str, Marker, "environments"),
+            extras=_get_sequence_as(d, str, _validate_normalized_name, "extras"),
+            dependency_groups=_get_sequence(d, str, "dependency-groups"),
+            default_groups=_get_sequence(d, str, "default-groups"),
+            created_by=_get_required(d, str, "created-by"),
+            requires_python=_get_as(d, str, SpecifierSet, "requires-python"),
+            packages=_get_required_sequence_of_objects(d, Package, "packages"),
+            tool=_get(d, Mapping, "tool"),  # type: ignore[type-abstract]
+        )
+        if not Version("1") <= pylock.lock_version < Version("2"):
+            raise PylockUnsupportedVersionError(
+                f"pylock version {pylock.lock_version} is not supported"
+            )
+        if pylock.lock_version > Version("1.0"):
+            _logger.warning(
+                "pylock minor version %s is not supported", pylock.lock_version
+            )
+        return pylock
+
+    @classmethod
+    def from_dict(cls, d: Mapping[str, Any], /) -> Self:
+        """Create and validate a Pylock instance from a TOML dictionary.
+
+        Raises :class:`PylockValidationError` if the input data is not
+        spec-compliant.
+        """
+        return cls._from_dict(d)
+
+    def to_dict(self) -> Mapping[str, Any]:
+        """Convert the Pylock instance to a TOML dictionary."""
+        return dataclasses.asdict(self, dict_factory=_toml_dict_factory)
+
+    def validate(self) -> None:
+        """Validate the Pylock instance against the specification.
+
+        Raises :class:`PylockValidationError` otherwise."""
+        self.from_dict(self.to_dict())
+
+    def select(
+        self,
+        *,
+        environment: Environment | None = None,
+        tags: Sequence[Tag] | None = None,
+        extras: Collection[str] | None = None,
+        dependency_groups: Collection[str] | None = None,
+    ) -> Iterator[
+        tuple[
+            Package,
+            PackageVcs
+            | PackageDirectory
+            | PackageArchive
+            | PackageWheel
+            | PackageSdist,
+        ]
+    ]:
+        """Select what to install from the lock file.
+
+        The *environment* and *tags* parameters represent the environment being
+        selected for. If unspecified, ``packaging.markers.default_environment()`` and
+        ``packaging.tags.sys_tags()`` are used.
+
+        The *extras* parameter represents the extras to install.
+
+        The *dependency_groups* parameter represents the groups to install. If
+        unspecified, the default groups are used.
+
+        This method must be used on valid Pylock instances (i.e. one obtained
+        from :meth:`Pylock.from_dict` or if constructed manually, after calling
+        :meth:`Pylock.validate`).
+        """
+        compatible_tags_selector = create_compatible_tags_selector(tags or sys_tags())
+
+        # #. Gather the extras and dependency groups to install and set ``extras`` and
+        #    ``dependency_groups`` for marker evaluation, respectively.
+        #
+        #    #. ``extras`` SHOULD be set to the empty set by default.
+        #    #. ``dependency_groups`` SHOULD be the set created from
+        #       :ref:`pylock-default-groups` by default.
+        env = cast(
+            "dict[str, str | frozenset[str]]",
+            dict(
+                environment or {},  # Marker.evaluate will fill-up
+                extras=frozenset(extras or []),
+                dependency_groups=frozenset(
+                    (self.default_groups or [])
+                    if dependency_groups is None  # to allow selecting no group
+                    else dependency_groups
+                ),
+            ),
+        )
+        env_python_full_version = (
+            environment["python_full_version"]
+            if environment
+            else default_environment()["python_full_version"]
+        )
+
+        # #. Check if the metadata version specified by :ref:`pylock-lock-version` is
+        #    supported; an error or warning MUST be raised as appropriate.
+        # Covered by lock.validate() which is a precondition for this method.
+
+        # #. If :ref:`pylock-requires-python` is specified, check that the environment
+        #    being installed for meets the requirement; an error MUST be raised if it is
+        #    not met.
+        if self.requires_python and not self.requires_python.contains(
+            env_python_full_version,
+        ):
+            raise PylockSelectError(
+                f"python_full_version {env_python_full_version!r} "
+                f"in provided environment does not satisfy the Python version "
+                f"requirement {str(self.requires_python)!r}"
+            )
+
+        # #. If :ref:`pylock-environments` is specified, check that at least one of the
+        #    environment marker expressions is satisfied; an error MUST be raised if no
+        #    expression is satisfied.
+        if self.environments:
+            for env_marker in self.environments:
+                if env_marker.evaluate(
+                    cast("dict[str, str]", environment or {}), context="requirement"
+                ):
+                    break
+            else:
+                raise PylockSelectError(
+                    "Provided environment does not satisfy any of the "
+                    "environments specified in the lock file"
+                )
+
+        # #. For each package listed in :ref:`pylock-packages`:
+        selected_packages_by_name: dict[str, tuple[int, Package]] = {}
+        for package_index, package in enumerate(self.packages):
+            # #. If :ref:`pylock-packages-marker` is specified, check if it is
+            #    satisfied;if it isn't, skip to the next package.
+            if package.marker and not package.marker.evaluate(env, context="lock_file"):
+                continue
+
+            # #. If :ref:`pylock-packages-requires-python` is specified, check if it is
+            #    satisfied; an error MUST be raised if it isn't.
+            if package.requires_python and not package.requires_python.contains(
+                env_python_full_version,
+            ):
+                raise PylockSelectError(
+                    f"python_full_version {env_python_full_version!r} "
+                    f"in provided environment does not satisfy the Python version "
+                    f"requirement {str(package.requires_python)!r} for package "
+                    f"{package.name!r} at packages[{package_index}]"
+                )
+
+            # #. Check that no other conflicting instance of the package has been slated
+            #    to be installed; an error about the ambiguity MUST be raised otherwise.
+            if package.name in selected_packages_by_name:
+                raise PylockSelectError(
+                    f"Multiple packages with the name {package.name!r} are "
+                    f"selected at packages[{package_index}] and "
+                    f"packages[{selected_packages_by_name[package.name][0]}]"
+                )
+
+            # #. Check that the source of the package is specified appropriately (i.e.
+            #    there are no conflicting sources in the package entry);
+            #    an error MUST be raised if any issues are found.
+            # Covered by lock.validate() which is a precondition for this method.
+
+            # #. Add the package to the set of packages to install.
+            selected_packages_by_name[package.name] = (package_index, package)
+
+        # #. For each package to be installed:
+        for package_index, package in selected_packages_by_name.values():
+            # - If :ref:`pylock-packages-vcs` is set:
+            if package.vcs is not None:
+                yield package, package.vcs
+
+            # - Else if :ref:`pylock-packages-directory` is set:
+            elif package.directory is not None:
+                yield package, package.directory
+
+            # - Else if :ref:`pylock-packages-archive` is set:
+            elif package.archive is not None:
+                yield package, package.archive
+
+            # - Else if there are entries for :ref:`pylock-packages-wheels`:
+            elif package.wheels:
+                # #. Look for the appropriate wheel file based on
+                #    :ref:`pylock-packages-wheels-name`; if one is not found then move
+                #    on to :ref:`pylock-packages-sdist` or an error MUST be raised about
+                #    a lack of source for the project.
+                best_wheel = next(
+                    compatible_tags_selector(
+                        (wheel, parse_wheel_filename(wheel.filename)[-1])
+                        for wheel in package.wheels
+                    ),
+                    None,
+                )
+                if best_wheel:
+                    yield package, best_wheel
+                elif package.sdist is not None:
+                    yield package, package.sdist
+                else:
+                    raise PylockSelectError(
+                        f"No wheel found matching the provided tags "
+                        f"for package {package.name!r} "
+                        f"at packages[{package_index}], "
+                        f"and no sdist available as a fallback"
+                    )
+
+            # - Else if no :ref:`pylock-packages-wheels` file is found or
+            #   :ref:`pylock-packages-sdist` is solely set:
+            elif package.sdist is not None:
+                yield package, package.sdist
+
+            else:
+                # Covered by lock.validate() which is a precondition for this method.
+                raise NotImplementedError  # pragma: no cover

LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/packaging/utils.py

@@ -0,0 +1,296 @@
+# This file is dual licensed under the terms of the Apache License, Version
+# 2.0, and the BSD License. See the LICENSE file in the root of this repository
+# for complete details.
+
+from __future__ import annotations
+
+import re
+from typing import NewType, Tuple, Union, cast
+
+from .tags import Tag, UnsortedTagsError, parse_tag
+from .version import InvalidVersion, Version, _TrimmedRelease
+
+__all__ = [
+    "BuildTag",
+    "InvalidName",
+    "InvalidSdistFilename",
+    "InvalidWheelFilename",
+    "NormalizedName",
+    "canonicalize_name",
+    "canonicalize_version",
+    "is_normalized_name",
+    "parse_sdist_filename",
+    "parse_wheel_filename",
+]
+
+
+def __dir__() -> list[str]:
+    return __all__
+
+
+BuildTag = Union[Tuple[()], Tuple[int, str]]
+
+NormalizedName = NewType("NormalizedName", str)
+"""
+A :class:`typing.NewType` of :class:`str`, representing a normalized name.
+"""
+
+
+class InvalidName(ValueError):
+    """
+    An invalid distribution name; users should refer to the packaging user guide.
+    """
+
+
+class InvalidWheelFilename(ValueError):
+    """
+    An invalid wheel filename was found, users should refer to PEP 427.
+    """
+
+
+class InvalidSdistFilename(ValueError):
+    """
+    An invalid sdist filename was found, users should refer to the packaging user guide.
+    """
+
+
+# Core metadata spec for `Name`
+_validate_regex = re.compile(
+    r"[a-z0-9]|[a-z0-9][a-z0-9._-]*[a-z0-9]", re.IGNORECASE | re.ASCII
+)
+_normalized_regex = re.compile(r"[a-z0-9]|[a-z0-9]([a-z0-9-](?!--))*[a-z0-9]", re.ASCII)
+# PEP 427: The build number must start with a digit.
+_build_tag_regex = re.compile(r"(\d+)(.*)", re.ASCII)
+
+
+def canonicalize_name(name: str, *, validate: bool = False) -> NormalizedName:
+    """
+    This function takes a valid Python package or extra name, and returns the
+    normalized form of it.
+
+    The return type is typed as :class:`NormalizedName`. This allows type
+    checkers to help require that a string has passed through this function
+    before use.
+
+    If **validate** is true, then the function will check if **name** is a valid
+    distribution name before normalizing.
+
+    :param str name: The name to normalize.
+    :param bool validate: Check whether the name is a valid distribution name.
+    :raises InvalidName: If **validate** is true and the name is not an
+        acceptable distribution name.
+
+    >>> from packaging.utils import canonicalize_name
+    >>> canonicalize_name("Django")
+    'django'
+    >>> canonicalize_name("oslo.concurrency")
+    'oslo-concurrency'
+    >>> canonicalize_name("requests")
+    'requests'
+    """
+    if validate and not _validate_regex.fullmatch(name):
+        raise InvalidName(f"name is invalid: {name!r}")
+    # Ensure all ``.`` and ``_`` are ``-``
+    # Emulates ``re.sub(r"[-_.]+", "-", name).lower()`` from PEP 503
+    # Much faster than re, and even faster than str.translate
+    value = name.lower().replace("_", "-").replace(".", "-")
+    # Condense repeats (faster than regex)
+    while "--" in value:
+        value = value.replace("--", "-")
+    return cast("NormalizedName", value)
+
+
+def is_normalized_name(name: str) -> bool:
+    """
+    Check if a name is already normalized (i.e. :func:`canonicalize_name` would
+    roundtrip to the same value).
+
+    :param str name: The name to check.
+
+    >>> from packaging.utils import is_normalized_name
+    >>> is_normalized_name("requests")
+    True
+    >>> is_normalized_name("Django")
+    False
+    """
+    return _normalized_regex.fullmatch(name) is not None
+
+
+def canonicalize_version(
+    version: Version | str, *, strip_trailing_zero: bool = True
+) -> str:
+    """Return a canonical form of a version as a string.
+
+    This function takes a string representing a package version (or a
+    :class:`~packaging.version.Version` instance), and returns the
+    normalized form of it. By default, it strips trailing zeros from
+    the release segment.
+
+    >>> from packaging.utils import canonicalize_version
+    >>> canonicalize_version('1.0.1')
+    '1.0.1'
+
+    Per PEP 625, versions may have multiple canonical forms, differing
+    only by trailing zeros.
+
+    >>> canonicalize_version('1.0.0')
+    '1'
+    >>> canonicalize_version('1.0.0', strip_trailing_zero=False)
+    '1.0.0'
+
+    Invalid versions are returned unaltered.
+
+    >>> canonicalize_version('foo bar baz')
+    'foo bar baz'
+
+    >>> canonicalize_version('1.4.0.0.0')
+    '1.4'
+    """
+    if isinstance(version, str):
+        try:
+            version = Version(version)
+        except InvalidVersion:
+            return str(version)
+    return str(_TrimmedRelease(version) if strip_trailing_zero else version)
+
+
+def parse_wheel_filename(
+    filename: str,
+    *,
+    validate_order: bool = False,
+) -> tuple[NormalizedName, Version, BuildTag, frozenset[Tag]]:
+    """
+    This function takes the filename of a wheel file, and parses it,
+    returning a tuple of name, version, build number, and tags.
+
+    The name part of the tuple is normalized and typed as
+    :class:`NormalizedName`. The version portion is an instance of
+    :class:`~packaging.version.Version`. The build number is ``()`` if
+    there is no build number in the wheel filename, otherwise a
+    two-item tuple of an integer for the leading digits and
+    a string for the rest of the build number. The tags portion is a
+    frozen set of :class:`~packaging.tags.Tag` instances (as the tag
+    string format allows multiple tags to be combined into a single
+    string).
+
+    If **validate_order** is true, compressed tag set components are
+    checked to be in sorted order as required by PEP 425.
+
+    :param str filename: The name of the wheel file.
+    :param bool validate_order: Check whether compressed tag set components
+        are in sorted order.
+    :raises InvalidWheelFilename: If the filename in question
+        does not follow the :ref:`wheel specification
+        <pypug:binary-distribution-format>`.
+
+    >>> from packaging.utils import parse_wheel_filename
+    >>> from packaging.tags import Tag
+    >>> from packaging.version import Version
+    >>> name, ver, build, tags = parse_wheel_filename("foo-1.0-py3-none-any.whl")
+    >>> name
+    'foo'
+    >>> ver == Version('1.0')
+    True
+    >>> tags == {Tag("py3", "none", "any")}
+    True
+    >>> not build
+    True
+
+    .. versionadded:: 26.1
+       The *validate_order* parameter.
+    """
+    if not filename.endswith(".whl"):
+        raise InvalidWheelFilename(
+            f"Invalid wheel filename (extension must be '.whl'): {filename!r}"
+        )
+
+    filename = filename[:-4]
+    dashes = filename.count("-")
+    if dashes not in (4, 5):
+        raise InvalidWheelFilename(
+            f"Invalid wheel filename (wrong number of parts): {filename!r}"
+        )
+
+    parts = filename.split("-", dashes - 2)
+    name_part = parts[0]
+    # See PEP 427 for the rules on escaping the project name.
+    if "__" in name_part or re.match(r"^[\w\d._]*$", name_part, re.UNICODE) is None:
+        raise InvalidWheelFilename(f"Invalid project name: {filename!r}")
+    name = canonicalize_name(name_part)
+
+    try:
+        version = Version(parts[1])
+    except InvalidVersion as e:
+        raise InvalidWheelFilename(
+            f"Invalid wheel filename (invalid version): {filename!r}"
+        ) from e
+
+    if dashes == 5:
+        build_part = parts[2]
+        build_match = _build_tag_regex.match(build_part)
+        if build_match is None:
+            raise InvalidWheelFilename(
+                f"Invalid build number: {build_part} in {filename!r}"
+            )
+        build = cast("BuildTag", (int(build_match.group(1)), build_match.group(2)))
+    else:
+        build = ()
+    tag_str = parts[-1]
+    try:
+        tags = parse_tag(tag_str, validate_order=validate_order)
+    except UnsortedTagsError:
+        raise InvalidWheelFilename(
+            f"Invalid wheel filename (compressed tag set components must be in "
+            f"sorted order per PEP 425): {filename!r}"
+        ) from None
+    return (name, version, build, tags)
+
+
+def parse_sdist_filename(filename: str) -> tuple[NormalizedName, Version]:
+    """
+    This function takes the filename of a sdist file (as specified
+    in the `Source distribution format`_ documentation), and parses
+    it, returning a tuple of the normalized name and version as
+    represented by an instance of :class:`~packaging.version.Version`.
+
+    :param str filename: The name of the sdist file.
+    :raises InvalidSdistFilename: If the filename does not end
+        with an sdist extension (``.zip`` or ``.tar.gz``), or if it does not
+        contain a dash separating the name and the version of the distribution.
+
+    >>> from packaging.utils import parse_sdist_filename
+    >>> from packaging.version import Version
+    >>> name, ver = parse_sdist_filename("foo-1.0.tar.gz")
+    >>> name
+    'foo'
+    >>> ver == Version('1.0')
+    True
+
+    .. _Source distribution format: https://packaging.python.org/specifications/source-distribution-format/#source-distribution-file-name
+    """
+    if filename.endswith(".tar.gz"):
+        file_stem = filename[: -len(".tar.gz")]
+    elif filename.endswith(".zip"):
+        file_stem = filename[: -len(".zip")]
+    else:
+        raise InvalidSdistFilename(
+            f"Invalid sdist filename (extension must be '.tar.gz' or '.zip'):"
+            f" {filename!r}"
+        )
+
+    # We are requiring a PEP 440 version, which cannot contain dashes,
+    # so we split on the last dash.
+    name_part, sep, version_part = file_stem.rpartition("-")
+    if not sep:
+        raise InvalidSdistFilename(f"Invalid sdist filename: {filename!r}")
+
+    name = canonicalize_name(name_part)
+
+    try:
+        version = Version(version_part)
+    except InvalidVersion as e:
+        raise InvalidSdistFilename(
+            f"Invalid sdist filename (invalid version): {filename!r}"
+        ) from e
+
+    return (name, version)

LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/cohere_asr/feature_extraction_cohere_asr.py

@@ -0,0 +1,374 @@
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import numpy as np
+import torch
+
+from ...feature_extraction_sequence_utils import SequenceFeatureExtractor
+from ...feature_extraction_utils import BatchFeature
+from ...utils import TensorType, is_librosa_available, logging
+from ...utils.import_utils import requires
+
+
+if is_librosa_available():
+    import librosa
+
+
+EPSILON = 1e-5
+LOG_ZERO_GUARD_VALUE = 2**-24
+
+
+logger = logging.get_logger(__name__)
+
+
+@requires(backends=("torch", "librosa"))
+class CohereAsrFeatureExtractor(SequenceFeatureExtractor):
+    r"""
+    Constructs a CohereAsr feature extractor.
+
+    This feature extractor inherits from [`~feature_extraction_sequence_utils.SequenceFeatureExtractor`] which contains
+    most of the main methods. Users should refer to this superclass for more information regarding those methods.
+
+    This class extracts mel-filter bank features from raw speech using a custom numpy implementation of the `Short Time
+    Fourier Transform` which should match pytorch's `torch.stft` equivalent.
+
+    Args:
+        feature_size (`int`, *optional*, defaults to 128):
+            The feature dimension of the extracted features.
+        sampling_rate (`int`, *optional*, defaults to 16000):
+            The sampling rate at which the audio files should be digitalized expressed in hertz (Hz).
+        hop_length (`int`, *optional*, defaults to 160):
+            Length of the overlapping windows for the STFT used to obtain the Mel Frequency coefficients.
+        n_fft (`int`, *optional*, defaults to 512):
+            Size of the Fourier transform.
+        win_length (`int`, *optional*, defaults to 400):
+            The window length for the STFT computation.
+        preemphasis (`float`, *optional*, defaults to 0.97):
+            A preemphasis filter coefficient. 0.0 means no preemphasis filter.
+        padding_value (`float`, *optional*, defaults to 0.0):
+            Padding value used to pad the audio. Should correspond to silences.
+        dither (`float`, *optional*, defaults to 1e-05):
+            Amount of deterministic dither noise to add before feature extraction. Each sample is seeded by its
+            valid waveform length so that dither is batch-composition invariant. Set to 0.0 to disable.
+        max_audio_clip_s (`float`, *optional*, defaults to 35.0):
+            Maximum duration in seconds for a single audio chunk. Audio longer than
+            `max_audio_clip_s - overlap_chunk_second` is split at energy-based boundaries.
+        overlap_chunk_second (`float`, *optional*, defaults to 5.0):
+            Size in seconds of the boundary search window used when splitting long audio. This is not actual
+            overlap between chunks — it defines how far back from the chunk boundary to search for a quiet
+            split point.
+        min_energy_window_samples (`int`, *optional*, defaults to 1600):
+            Size in samples of the sliding window used to find the quietest point when splitting audio chunks.
+    """
+
+    model_input_names = ["input_features", "attention_mask"]
+
+    def __init__(
+        self,
+        feature_size=128,
+        sampling_rate=16000,
+        hop_length=160,
+        n_fft=512,
+        win_length=400,
+        preemphasis=0.97,
+        padding_value=0.0,
+        dither=1e-5,
+        max_audio_clip_s=35.0,
+        overlap_chunk_second=5.0,
+        min_energy_window_samples=1600,
+        **kwargs,
+    ):
+        super().__init__(feature_size=feature_size, sampling_rate=sampling_rate, padding_value=padding_value, **kwargs)
+
+        self.hop_length = hop_length
+        self.n_fft = n_fft
+        self.win_length = win_length
+        self.preemphasis = preemphasis
+        self.dither = dither
+        self.max_audio_clip_s = max_audio_clip_s
+        self.overlap_chunk_second = overlap_chunk_second
+        self.min_energy_window_samples = min_energy_window_samples
+
+        # TODO: @eustlb, for now we use librosa to compute the mel filters
+        # indeed mel_filter_bank uses np.float64 (while librosa uses np.float32), giving numerical differences
+        mel_filters = librosa.filters.mel(
+            sr=sampling_rate, n_fft=n_fft, n_mels=feature_size, fmin=0.0, fmax=sampling_rate / 2, norm="slaney"
+        )
+        self.mel_filters = torch.from_numpy(mel_filters).to(torch.float32)
+
+    def _find_split_point_energy(self, waveform: torch.Tensor, start_idx: int, end_idx: int) -> int:
+        segment = waveform[start_idx:end_idx]
+        if segment.shape[0] <= self.min_energy_window_samples:
+            return (start_idx + end_idx) // 2
+
+        min_energy = float("inf")
+        quietest_idx = start_idx
+        upper = segment.shape[0] - self.min_energy_window_samples
+        for i in range(0, upper, self.min_energy_window_samples):
+            window = segment[i : i + self.min_energy_window_samples]
+            energy = torch.sqrt(torch.mean(window * window)).item()
+            if energy < min_energy:
+                min_energy = energy
+                quietest_idx = start_idx + i
+        return quietest_idx
+
+    def _split_audio_chunks_energy(self, waveform: torch.Tensor) -> list[torch.Tensor]:
+        chunk_size = max(1, int(round(self.max_audio_clip_s * self.sampling_rate)))
+        boundary_context_size = max(1, int(round(self.overlap_chunk_second * self.sampling_rate)))
+        total_samples = waveform.shape[0]
+
+        if total_samples <= chunk_size:
+            return [waveform]
+
+        chunks_meta: list[tuple[int, int]] = []
+        idx = 0
+        while idx < total_samples:
+            if idx + chunk_size >= total_samples:
+                chunks_meta.append((idx, total_samples))
+                break
+
+            search_start = max(idx, idx + chunk_size - boundary_context_size)
+            search_end = min(idx + chunk_size, total_samples)
+            if search_end <= search_start:
+                split_point = idx + chunk_size
+            else:
+                split_point = self._find_split_point_energy(waveform, search_start, search_end)
+
+            split_point = max(idx + 1, min(split_point, total_samples))
+            chunks_meta.append((idx, split_point))
+            idx = split_point
+
+        return [waveform[start:end] for start, end in chunks_meta if end > start]
+
+    def _apply_dither(self, waveform: torch.Tensor, audio_lengths: torch.Tensor) -> torch.Tensor:
+        if self.dither <= 0:
+            return waveform
+        generator = torch.Generator(device=waveform.device)
+        for i in range(waveform.shape[0]):
+            valid_samples = min(int(audio_lengths[i].item()), waveform.shape[1])
+            if valid_samples <= 0:
+                continue
+            generator.manual_seed(valid_samples)
+            noise = torch.randn(valid_samples, dtype=waveform.dtype, device=waveform.device, generator=generator)
+            waveform[i, :valid_samples] += self.dither * noise
+        return waveform
+
+    def _torch_extract_fbank_features(self, waveform, device="cpu"):
+        # spectrogram
+        window = torch.hann_window(self.win_length, periodic=False, device=device)
+        stft = torch.stft(
+            waveform,
+            self.n_fft,
+            hop_length=self.hop_length,
+            win_length=self.win_length,
+            window=window,
+            return_complex=True,
+            pad_mode="constant",
+        )
+        # Let's match original implementation
+        magnitudes = torch.view_as_real(stft)
+        magnitudes = torch.sqrt(magnitudes.pow(2).sum(-1))
+        magnitudes = magnitudes.pow(2)
+
+        # log mel spectrogram
+        mel_filters = self.mel_filters.to(device)
+        mel_spec = mel_filters @ magnitudes
+        mel_spec = torch.log(mel_spec + LOG_ZERO_GUARD_VALUE)
+
+        # (batch_size, num_mel_filters, num_frames) -> (batch_size, num_frames, num_mel_filters)
+        mel_spec = mel_spec.permute(0, 2, 1)
+
+        return mel_spec
+
+    def __call__(
+        self,
+        raw_speech: np.ndarray | list[float] | list[np.ndarray] | list[list[float]],
+        truncation: bool = False,
+        pad_to_multiple_of: int | None = None,
+        return_tensors: str | TensorType | None = None,
+        return_attention_mask: bool | None = None,
+        padding: str | None = "longest",
+        max_length: int | None = None,
+        sampling_rate: int | None = None,
+        do_normalize: bool | None = None,
+        device: str | None = "cpu",
+        return_token_timestamps: bool | None = None,
+        **kwargs,
+    ) -> BatchFeature:
+        """
+        Main method to featurize and prepare for the model one or several sequence(s). Implementation uses PyTorch for
+        the STFT computation if available, otherwise a slower NumPy based one.
+
+        Args:
+            raw_speech (`np.ndarray`, `list[float]`, `list[np.ndarray]`, `list[list[float]]`):
+                The sequence or batch of sequences to be padded. Each sequence can be a numpy array, a list of float
+                values, a list of numpy arrays or a list of list of float values. Must be mono channel audio, not
+                stereo, i.e. single float per timestep.
+            truncation (`bool`, *optional*, default to `True`):
+                Activates truncation to cut input sequences longer than *max_length* to *max_length*.
+            pad_to_multiple_of (`int`, *optional*, defaults to None):
+                If set will pad the sequence to a multiple of the provided value.
+
+                This is especially useful to enable the use of Tensor Cores on NVIDIA hardware with compute capability
+                `>= 7.5` (Volta), or on TPUs which benefit from having sequence lengths be a multiple of 128.
+            return_attention_mask (`bool`, *optional*):
+                Whether to return the attention mask. If left to the default, will return the attention mask according
+                to the specific feature_extractor's default.
+
+                [What are attention masks?](../glossary#attention-mask)
+
+                <Tip>
+
+                For CohereAsr models, `attention_mask` should always be passed for batched inference, to avoid subtle
+                bugs.
+
+                </Tip>
+
+            return_tensors (`str` or [`~utils.TensorType`], *optional*):
+                If set, will return tensors instead of list of python integers. Acceptable values are:
+
+                - `'tf'`: Return TensorFlow `tf.constant` objects.
+                - `'pt'`: Return PyTorch `torch.Tensor` objects.
+                - `'np'`: Return Numpy `np.ndarray` objects.
+            sampling_rate (`int`, *optional*):
+                The sampling rate at which the `raw_speech` input was sampled. It is strongly recommended to pass
+                `sampling_rate` at the forward call to prevent silent errors and allow automatic speech recognition
+                pipeline.
+            padding_value (`float`, *optional*, defaults to 0.0):
+                The value that is used to fill the padding values / vectors.
+            do_normalize (`bool`, *optional*, defaults to `False`):
+                Whether or not to zero-mean unit-variance normalize the input. Normalizing can help to significantly
+                improve the performance of the model.
+            device (`str`, *optional*, defaults to `'cpu'`):
+                Specifies the device for computation of the log-mel spectrogram of audio signals in the
+                `_torch_extract_fbank_features` method. (e.g., "cpu", "cuda")
+            return_token_timestamps (`bool`, *optional*, defaults to `None`):
+                Deprecated. Use `return_attention_mask` instead from which the number of frames can be inferred.
+
+                Whether or not to return the number of frames of the input raw_speech.
+                These num_frames can be used by the model to compute word level timestamps.
+        """
+        if sampling_rate is not None:
+            if sampling_rate != self.sampling_rate:
+                raise ValueError(
+                    f"The model corresponding to this feature extractor: {self.__class__.__name__} was trained using a"
+                    f" sampling rate of {self.sampling_rate}. Please make sure that the provided `raw_speech` input"
+                    f" was sampled with {self.sampling_rate} and not {sampling_rate}."
+                )
+        else:
+            logger.warning(
+                f"It is strongly recommended to pass the `sampling_rate` argument to `{self.__class__.__name__}()`. "
+                "Failing to do so can result in silent errors that might be hard to debug."
+            )
+
+        # Convert to torch tensor
+        if isinstance(raw_speech, np.ndarray):
+            raw_speech = torch.tensor(raw_speech)
+        elif isinstance(raw_speech, (list, tuple)) and isinstance(raw_speech[0], np.ndarray):
+            raw_speech = [torch.tensor(speech) for speech in raw_speech]
+
+        is_batched_torch = isinstance(raw_speech, torch.Tensor) and len(raw_speech.shape) > 1
+        if is_batched_torch and len(raw_speech.shape) > 2:
+            logger.warning(
+                f"Only mono-channel audio is supported for input to {self.__class__.__name__}. "
+                "We will take the mean of the channels to convert to mono."
+            )
+            raw_speech = raw_speech.mean(-1)
+
+        is_batched_sequence = isinstance(raw_speech, (list, tuple))
+        if is_batched_sequence:
+            for speech in raw_speech:
+                if len(speech.shape) > 1:
+                    logger.warning(
+                        f"Only mono-channel audio is supported for input to {self.__class__.__name__}. "
+                        "We will take the mean of the channels to convert to mono."
+                    )
+                    speech = speech.mean(-1)
+
+        if is_batched_torch or is_batched_sequence:
+            raw_speech = [speech.to(torch.float32) for speech in raw_speech]
+        else:
+            raw_speech = [raw_speech.to(torch.float32)]
+
+        # Chunk long audio at energy-based boundaries
+        fast_path_threshold_s = max(0.0, self.max_audio_clip_s - self.overlap_chunk_second)
+        audio_chunk_index: list[tuple[int, int | None]] = []
+        chunked_speech: list[torch.Tensor] = []
+        for sample_idx, speech in enumerate(raw_speech):
+            duration_s = speech.shape[0] / self.sampling_rate
+            if duration_s <= fast_path_threshold_s:
+                chunked_speech.append(speech)
+                audio_chunk_index.append((sample_idx, None))
+            else:
+                chunks = self._split_audio_chunks_energy(speech)
+                for chunk_idx, chunk in enumerate(chunks):
+                    chunked_speech.append(chunk)
+                    audio_chunk_index.append((sample_idx, chunk_idx))
+
+        raw_speech = [speech[:, None] for speech in chunked_speech]
+
+        audio_lengths = [len(speech) for speech in raw_speech]
+        batched_speech = BatchFeature({"input_features": raw_speech, "audio_lengths": audio_lengths})
+
+        padded_inputs = self.pad(
+            batched_speech,
+            padding=padding,
+            max_length=max_length,
+            truncation=truncation,
+            pad_to_multiple_of=pad_to_multiple_of,
+            return_tensors="pt",
+        )
+        input_features = padded_inputs.input_features.squeeze(-1)
+
+        # dithering
+        input_features = self._apply_dither(input_features, padded_inputs.audio_lengths)
+
+        # preemphasis
+        if self.preemphasis is not None:
+            timemask = torch.arange(input_features.shape[1], device=input_features.device).unsqueeze(
+                0
+            ) < padded_inputs.audio_lengths.unsqueeze(1)
+            input_features = torch.cat(
+                [input_features[:, :1], input_features[:, 1:] - self.preemphasis * input_features[:, :-1]], dim=1
+            )
+            input_features = input_features.masked_fill(~timemask, 0.0)
+
+        input_features = self._torch_extract_fbank_features(input_features, device)
+        features_lengths = torch.floor_divide(
+            padded_inputs.audio_lengths + self.n_fft // 2 * 2 - self.n_fft, self.hop_length
+        )
+        attention_mask = torch.arange(input_features.shape[1], device=device)[None, :] < features_lengths[:, None]
+
+        # normalize mel features, ignoring padding
+        mask = attention_mask.unsqueeze(-1)
+        input_features_masked = input_features * mask
+        mean = input_features_masked.sum(dim=1) / features_lengths.unsqueeze(-1)
+        mean = mean.unsqueeze(1)
+        variance = ((input_features_masked - mean) ** 2 * mask).sum(dim=1) / (features_lengths - 1).unsqueeze(-1)
+        std = torch.sqrt(variance).unsqueeze(1)
+        input_features = (input_features - mean) / (std + EPSILON)
+        input_features *= mask
+
+        result = BatchFeature(
+            data={
+                "input_features": input_features,
+                "attention_mask": attention_mask,
+            },
+            tensor_type=return_tensors,
+        )
+        result["audio_chunk_index"] = audio_chunk_index
+        return result
+
+
+__all__ = ["CohereAsrFeatureExtractor"]

LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/cohere_asr/processing_cohere_asr.py

@@ -0,0 +1,188 @@
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from ...audio_utils import AudioInput, make_list_of_audio
+from ...processing_utils import ProcessingKwargs, ProcessorMixin, Unpack
+from ...tokenization_utils_base import PreTokenizedInput, TextInput
+from ...utils import auto_docstring, is_torch_available, logging
+from ...utils.import_utils import requires
+
+
+if is_torch_available():
+    import torch
+
+
+LANGUAGES = {"ar", "de", "el", "en", "es", "fr", "it", "ja", "ko", "nl", "pl", "pt", "vi", "zh"}
+_NO_SPACE_LANGS = {"ja", "zh"}
+
+
+logger = logging.get_logger(__name__)
+
+
+class CohereAsrProcessorKwargs(ProcessingKwargs, total=False):
+    _defaults = {
+        "audio_kwargs": {
+            "sampling_rate": 16000,
+            "padding": "longest",
+            "return_attention_mask": True,
+        },
+        "text_kwargs": {
+            "padding": True,
+            "padding_side": "right",
+            "add_special_tokens": False,
+        },
+        "common_kwargs": {"return_tensors": "pt"},
+    }
+
+
+@auto_docstring
+@requires(backends=("torch",))
+class CohereAsrProcessor(ProcessorMixin):
+    def __init__(self, feature_extractor, tokenizer):
+        super().__init__(feature_extractor, tokenizer)
+
+    def get_decoder_prompt_ids(self, language: str, punctuation: bool = True) -> list[int]:
+        """Build the decoder prompt token IDs for the given language and punctuation settings."""
+        if language not in LANGUAGES:
+            raise ValueError(
+                f"Unsupported language: {language!r}. Supported languages: {', '.join(sorted(LANGUAGES))}."
+            )
+        pnc_token = "<|pnc|>" if punctuation else "<|nopnc|>"
+        tokens = [
+            "▁",
+            "<|startofcontext|>",
+            "<|startoftranscript|>",
+            "<|emo:undefined|>",
+            f"<|{language}|>",
+            f"<|{language}|>",
+            pnc_token,
+            "<|noitn|>",
+            "<|notimestamp|>",
+            "<|nodiarize|>",
+        ]
+        return self.tokenizer.convert_tokens_to_ids(tokens)
+
+    @auto_docstring
+    def __call__(
+        self,
+        audio: AudioInput,
+        language: str,
+        text: TextInput | PreTokenizedInput | list[TextInput] | list[PreTokenizedInput] | None = None,
+        punctuation: bool = True,
+        sampling_rate: int | None = None,
+        **kwargs: Unpack[CohereAsrProcessorKwargs],
+    ):
+        r"""
+        language (`str`):
+            Language code (e.g. `"en"`, `"es"`, `"fr"`) used to build the decoder prompt. The processor
+            constructs the full decoder prompt and returns `decoder_input_ids` alongside the audio features.
+        punctuation (`bool`, defaults to `True`):
+            Whether to enable punctuation in the decoder prompt.
+        sampling_rate (`int`, *optional*):
+            The sampling rate of the input audio in Hz. This should match the sampling rate expected by the feature
+            extractor (defaults to 16000 Hz). If provided, it will be validated against the processor's expected
+            sampling rate, and an error will be raised if they don't match. If not provided, a warning will be
+            issued and the default sampling rate will be assumed.
+        """
+        audio = make_list_of_audio(audio)
+
+        output_kwargs = self._merge_kwargs(
+            CohereAsrProcessorKwargs,
+            tokenizer_init_kwargs=self.tokenizer.init_kwargs,
+            **kwargs,
+        )
+
+        if sampling_rate is None:
+            logger.warning_once(
+                f"You've provided audio without specifying the sampling rate. It will be assumed to be {output_kwargs['audio_kwargs']['sampling_rate']}, which can result in silent errors."
+            )
+        elif sampling_rate != output_kwargs["audio_kwargs"]["sampling_rate"]:
+            raise ValueError(
+                f"The sampling rate of the audio ({sampling_rate}) does not match the sampling rate of the processor ({output_kwargs['audio_kwargs']['sampling_rate']}). Please provide resampled the audio to the expected sampling rate."
+            )
+
+        inputs = self.feature_extractor(audio, **output_kwargs["audio_kwargs"])
+
+        prompt_ids = self.get_decoder_prompt_ids(language=language, punctuation=punctuation)
+        batch_size = inputs["input_features"].shape[0]
+        inputs["decoder_input_ids"] = torch.tensor([prompt_ids] * batch_size, dtype=torch.long)
+
+        if text is not None:
+            encodings = self.tokenizer(text, **output_kwargs["text_kwargs"])
+            inputs["labels"] = encodings["input_ids"]
+
+        return inputs
+
+    def decode(self, *args, audio_chunk_index=None, language=None, **kwargs):
+        texts = self.tokenizer.decode(*args, **kwargs)
+        if audio_chunk_index is None:
+            return texts
+        if language is None:
+            raise ValueError("`language` must be provided when `audio_chunk_index` is given.")
+        separator = "" if language in _NO_SPACE_LANGS else " "
+        return self._reassemble_chunk_texts(texts, audio_chunk_index, separator)
+
+    @staticmethod
+    def _reassemble_chunk_texts(
+        texts: list[str],
+        audio_chunk_index: list[tuple[int, int | None]],
+        separator: str = " ",
+    ) -> list[str]:
+        """Reassemble per-chunk transcription texts back into per-sample strings.
+
+        When audio inputs are longer than the feature extractor's `max_audio_clip_s`, they are split into
+        overlapping chunks before being fed to the model. This means a single original audio sample can
+        produce multiple decoded text segments. This method reverses that chunking: it groups the decoded
+        texts by their original sample index using `chunk_map`, orders the chunks, and joins them
+        with `separator` to reconstruct one transcription string per input sample.
+
+        Args:
+            texts: Decoded text strings, one per model output (i.e. one per chunk).
+            audio_chunk_index: List of `(sample_idx, chunk_idx)` tuples that map each entry in
+                `texts` back to its original sample and chunk position. A `chunk_idx` of `None`
+                indicates the sample was not chunked.
+            separator: String used to join chunks belonging to the same sample. Defaults to a
+                space; callers pass an empty string for languages that don't use spaces between
+                words (e.g. Chinese, Japanese).
+
+        Returns:
+            A list of reassembled transcription strings, one per original input sample.
+        """
+        max_sample_idx = max(sample_idx for sample_idx, _ in audio_chunk_index)
+        outputs = [""] * (max_sample_idx + 1)
+        chunked = {}
+
+        for (sample_idx, chunk_idx), text in zip(audio_chunk_index, texts):
+            if chunk_idx is None:
+                outputs[sample_idx] = text
+            else:
+                if sample_idx not in chunked:
+                    chunked[sample_idx] = []
+                chunked[sample_idx].append((chunk_idx, text))
+
+        for sample_idx, chunk_items in chunked.items():
+            chunk_items.sort(key=lambda item: item[0])
+            non_empty = [t for _, t in chunk_items if t and t.strip()]
+            parts = [non_empty[0].rstrip()] + [t.strip() for t in non_empty[1:]]
+            outputs[sample_idx] = separator.join(parts)
+
+        return outputs
+
+    @property
+    def model_input_names(self):
+        feature_extractor_input_names = self.feature_extractor.model_input_names
+        return feature_extractor_input_names + ["labels"]
+
+
+__all__ = ["CohereAsrProcessor"]

LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/esm/openfold_utils/chunk_utils.py

@@ -0,0 +1,398 @@
+# Copyright 2021 AlQuraishi Laboratory
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import logging
+import math
+from collections.abc import Callable, Iterable, Sequence
+from functools import partial
+from typing import Any
+
+import torch
+
+from .tensor_utils import tensor_tree_map, tree_map
+
+
+def _fetch_dims(tree: dict | list | tuple | torch.Tensor) -> list[tuple[int, ...]]:
+    shapes = []
+    if isinstance(tree, dict):
+        for v in tree.values():
+            shapes.extend(_fetch_dims(v))
+    elif isinstance(tree, (list, tuple)):
+        for t in tree:
+            shapes.extend(_fetch_dims(t))
+    elif isinstance(tree, torch.Tensor):
+        shapes.append(tree.shape)
+    else:
+        raise TypeError("Not supported")
+
+    return shapes
+
+
+@torch.jit.ignore
+def _flat_idx_to_idx(flat_idx: int, dims: tuple[int, ...]) -> tuple[int, ...]:
+    idx = []
+    for d in reversed(dims):
+        idx.append(flat_idx % d)
+        flat_idx = flat_idx // d
+
+    return tuple(reversed(idx))
+
+
+@torch.jit.ignore
+def _get_minimal_slice_set(
+    start: Sequence[int],
+    end: Sequence[int],
+    dims: Sequence[int],
+    start_edges: Sequence[bool] | None = None,
+    end_edges: Sequence[bool] | None = None,
+) -> list[tuple[slice, ...]]:
+    """
+    Produces an ordered sequence of tensor slices that, when used in sequence on a tensor with shape dims, yields
+    tensors that contain every leaf in the contiguous range [start, end]. Care is taken to yield a short sequence of
+    slices, and perhaps even the shortest possible (I'm pretty sure it's the latter).
+
+    end is INCLUSIVE.
+    """
+
+    # start_edges and end_edges both indicate whether, starting from any given
+    # dimension, the start/end index is at the top/bottom edge of the
+    # corresponding tensor, modeled as a tree
+    def reduce_edge_list(l: list[bool]) -> None:
+        tally = True
+        for i in range(len(l)):
+            reversed_idx = -1 * (i + 1)
+            l[reversed_idx] &= tally
+            tally = l[reversed_idx]
+
+    if start_edges is None:
+        start_edges = [s == 0 for s in start]
+        reduce_edge_list(start_edges)
+    if end_edges is None:
+        end_edges = [e == (d - 1) for e, d in zip(end, dims)]
+        reduce_edge_list(end_edges)
+
+    # Base cases. Either start/end are empty and we're done, or the final,
+    # one-dimensional tensor can be simply sliced
+    if len(start) == 0:
+        return [()]
+    elif len(start) == 1:
+        return [(slice(start[0], end[0] + 1),)]
+
+    slices: list[tuple[slice, ...]] = []
+    path_list: list[slice] = []
+
+    # Dimensions common to start and end can be selected directly
+    for s, e in zip(start, end):
+        if s == e:
+            path_list.append(slice(s, s + 1))
+        else:
+            break
+
+    path: tuple[slice, ...] = tuple(path_list)
+    divergence_idx = len(path)
+
+    # start == end, and we're done
+    if divergence_idx == len(dims):
+        return [path]
+
+    def upper() -> tuple[tuple[slice, ...], ...]:
+        assert start_edges is not None
+        assert end_edges is not None
+
+        sdi = start[divergence_idx]
+        return tuple(
+            path + (slice(sdi, sdi + 1),) + s
+            for s in _get_minimal_slice_set(
+                start[divergence_idx + 1 :],
+                [d - 1 for d in dims[divergence_idx + 1 :]],
+                dims[divergence_idx + 1 :],
+                start_edges=start_edges[divergence_idx + 1 :],
+                end_edges=[True for _ in end_edges[divergence_idx + 1 :]],
+            )
+        )
+
+    def lower() -> tuple[tuple[slice, ...], ...]:
+        assert start_edges is not None
+        assert end_edges is not None
+
+        edi = end[divergence_idx]
+        return tuple(
+            path + (slice(edi, edi + 1),) + s
+            for s in _get_minimal_slice_set(
+                [0 for _ in start[divergence_idx + 1 :]],
+                end[divergence_idx + 1 :],
+                dims[divergence_idx + 1 :],
+                start_edges=[True for _ in start_edges[divergence_idx + 1 :]],
+                end_edges=end_edges[divergence_idx + 1 :],
+            )
+        )
+
+    # If both start and end are at the edges of the subtree rooted at
+    # divergence_idx, we can just select the whole subtree at once
+    if start_edges[divergence_idx] and end_edges[divergence_idx]:
+        slices.append(path + (slice(start[divergence_idx], end[divergence_idx] + 1),))
+    # If just start is at the edge, we can grab almost all of the subtree,
+    # treating only the ragged bottom edge as an edge case
+    elif start_edges[divergence_idx]:
+        slices.append(path + (slice(start[divergence_idx], end[divergence_idx]),))
+        slices.extend(lower())
+    # Analogous to the previous case, but the top is ragged this time
+    elif end_edges[divergence_idx]:
+        slices.extend(upper())
+        slices.append(path + (slice(start[divergence_idx] + 1, end[divergence_idx] + 1),))
+    # If both sides of the range are ragged, we need to handle both sides
+    # separately. If there's contiguous meat in between them, we can index it
+    # in one big chunk
+    else:
+        slices.extend(upper())
+        middle_ground = end[divergence_idx] - start[divergence_idx]
+        if middle_ground > 1:
+            slices.append(path + (slice(start[divergence_idx] + 1, end[divergence_idx]),))
+        slices.extend(lower())
+
+    return slices
+
+
+@torch.jit.ignore
+def _chunk_slice(t: torch.Tensor, flat_start: int, flat_end: int, no_batch_dims: int) -> torch.Tensor:
+    """
+    Equivalent to
+
+        t.reshape((-1,) + t.shape[no_batch_dims:])[flat_start:flat_end]
+
+    but without the need for the initial reshape call, which can be memory-intensive in certain situations. The only
+    reshape operations in this function are performed on sub-tensors that scale with (flat_end - flat_start), the chunk
+    size.
+    """
+
+    batch_dims = t.shape[:no_batch_dims]
+    start_idx = list(_flat_idx_to_idx(flat_start, batch_dims))
+    # _get_minimal_slice_set is inclusive
+    end_idx = list(_flat_idx_to_idx(flat_end - 1, batch_dims))
+
+    # Get an ordered list of slices to perform
+    slices = _get_minimal_slice_set(
+        start_idx,
+        end_idx,
+        batch_dims,
+    )
+
+    sliced_tensors = [t[s] for s in slices]
+
+    return torch.cat([s.view((-1,) + t.shape[no_batch_dims:]) for s in sliced_tensors])
+
+
+def chunk_layer(
+    layer: Callable,
+    inputs: dict[str, Any],
+    chunk_size: int,
+    no_batch_dims: int,
+    low_mem: bool = False,
+    _out: Any = None,
+    _add_into_out: bool = False,
+) -> Any:
+    """
+    Implements the "chunking" procedure described in section 1.11.8.
+
+    Layer outputs and inputs are assumed to be simple "pytrees," consisting only of (arbitrarily nested) lists, tuples,
+    and dicts with torch.Tensor leaves.
+
+    Args:
+        layer:
+            The layer to be applied chunk-wise
+        inputs:
+            A (non-nested) dictionary of keyworded inputs. All leaves must be tensors and must share the same batch
+            dimensions.
+        chunk_size:
+            The number of sub-batches per chunk. If multiple batch dimensions are specified, a "sub-batch" is defined
+            as a single indexing of all batch dimensions simultaneously (s.t. the number of sub-batches is the product
+            of the batch dimensions).
+        no_batch_dims:
+            How many of the initial dimensions of each input tensor can be considered batch dimensions.
+        low_mem:
+            Avoids flattening potentially large input tensors. Unnecessary in most cases, and is ever so slightly
+            slower than the default setting.
+    Returns:
+        The reassembled output of the layer on the inputs.
+    """
+    if not (len(inputs) > 0):
+        raise ValueError("Must provide at least one input")
+
+    initial_dims = [shape[:no_batch_dims] for shape in _fetch_dims(inputs)]
+    orig_batch_dims = tuple(max(s) for s in zip(*initial_dims))
+
+    def _prep_inputs(t: torch.Tensor) -> torch.Tensor:
+        if not low_mem:
+            if sum(t.shape[:no_batch_dims]) != no_batch_dims:
+                t = t.expand(orig_batch_dims + t.shape[no_batch_dims:])
+            t = t.reshape(-1, *t.shape[no_batch_dims:])
+        else:
+            t = t.expand(orig_batch_dims + t.shape[no_batch_dims:])
+        return t
+
+    prepped_inputs: dict[str, Any] = tensor_tree_map(_prep_inputs, inputs)
+    prepped_outputs = None
+    if _out is not None:
+        prepped_outputs = tensor_tree_map(lambda t: t.view([-1] + list(t.shape[no_batch_dims:])), _out)
+
+    flat_batch_dim = 1
+    for d in orig_batch_dims:
+        flat_batch_dim *= d
+
+    no_chunks = flat_batch_dim // chunk_size + (flat_batch_dim % chunk_size != 0)
+
+    def _select_chunk(t: torch.Tensor) -> torch.Tensor:
+        return t[i : i + chunk_size] if t.shape[0] != 1 else t
+
+    i = 0
+    out = prepped_outputs
+    for _ in range(no_chunks):
+        # Chunk the input
+        if not low_mem:
+            select_chunk = _select_chunk
+        else:
+            select_chunk = partial(
+                _chunk_slice,
+                flat_start=i,
+                flat_end=min(flat_batch_dim, i + chunk_size),
+                no_batch_dims=len(orig_batch_dims),
+            )
+
+        chunks: dict[str, Any] = tensor_tree_map(select_chunk, prepped_inputs)
+
+        # Run the layer on the chunk
+        output_chunk = layer(**chunks)
+
+        # Allocate space for the output
+        if out is None:
+            out = tensor_tree_map(lambda t: t.new_zeros((flat_batch_dim,) + t.shape[1:]), output_chunk)
+
+        # Put the chunk in its pre-allocated space
+        if isinstance(output_chunk, dict):
+
+            def assign(d1: dict, d2: dict) -> None:
+                for k, v in d1.items():
+                    if isinstance(v, dict):
+                        assign(v, d2[k])
+                    else:
+                        if _add_into_out:
+                            v[i : i + chunk_size] += d2[k]
+                        else:
+                            v[i : i + chunk_size] = d2[k]
+
+            assign(out, output_chunk)
+        elif isinstance(output_chunk, tuple):
+            for x1, x2 in zip(out, output_chunk):
+                if _add_into_out:
+                    x1[i : i + chunk_size] += x2
+                else:
+                    x1[i : i + chunk_size] = x2
+        elif isinstance(output_chunk, torch.Tensor):
+            if _add_into_out:
+                out[i : i + chunk_size] += output_chunk
+            else:
+                out[i : i + chunk_size] = output_chunk
+        else:
+            raise TypeError("Not supported")
+
+        i += chunk_size
+
+    out = tensor_tree_map(lambda t: t.view(orig_batch_dims + t.shape[1:]), out)
+
+    return out
+
+
+class ChunkSizeTuner:
+    def __init__(
+        self,
+        # Heuristically, runtimes for most of the modules in the network
+        # plateau earlier than this on all GPUs I've run the model on.
+        max_chunk_size: int = 512,
+    ):
+        self.max_chunk_size = max_chunk_size
+        self.cached_chunk_size: int | None = None
+        self.cached_arg_data: tuple | None = None
+
+    def _determine_favorable_chunk_size(self, fn: Callable, args: tuple, min_chunk_size: int) -> int:
+        logging.info("Tuning chunk size...")
+
+        if min_chunk_size >= self.max_chunk_size:
+            return min_chunk_size
+
+        candidates: list[int] = [2**l for l in range(int(math.log2(self.max_chunk_size)) + 1)]
+        candidates = [c for c in candidates if c > min_chunk_size]
+        candidates = [min_chunk_size] + candidates
+        candidates[-1] += 4
+
+        def test_chunk_size(chunk_size: int) -> bool:
+            try:
+                with torch.no_grad():
+                    fn(*args, chunk_size=chunk_size)
+                return True
+            except RuntimeError:
+                return False
+
+        min_viable_chunk_size_index = 0
+        i = len(candidates) - 1
+        while i > min_viable_chunk_size_index:
+            viable = test_chunk_size(candidates[i])
+            if not viable:
+                i = (min_viable_chunk_size_index + i) // 2
+            else:
+                min_viable_chunk_size_index = i
+                i = (i + len(candidates) - 1) // 2
+
+        return candidates[min_viable_chunk_size_index]
+
+    def _compare_arg_caches(self, ac1: Iterable, ac2: Iterable) -> bool:
+        consistent = True
+        for a1, a2 in zip(ac1, ac2):
+            assert type(ac1) is type(ac2)
+            if isinstance(ac1, (list, tuple)):
+                consistent &= self._compare_arg_caches(a1, a2)
+            elif isinstance(ac1, dict):
+                a1_items = [v for _, v in sorted(a1.items(), key=lambda x: x[0])]
+                a2_items = [v for _, v in sorted(a2.items(), key=lambda x: x[0])]
+                consistent &= self._compare_arg_caches(a1_items, a2_items)
+            else:
+                consistent &= a1 == a2
+
+        return consistent
+
+    def tune_chunk_size(
+        self,
+        representative_fn: Callable,
+        args: tuple,
+        min_chunk_size: int,
+    ) -> int:
+        consistent = True
+        arg_data: tuple = tree_map(lambda a: a.shape if isinstance(a, torch.Tensor) else a, args, object)
+        if self.cached_arg_data is not None:
+            # If args have changed shape/value, we need to re-tune
+            assert len(self.cached_arg_data) == len(arg_data)
+            consistent = self._compare_arg_caches(self.cached_arg_data, arg_data)
+        else:
+            # Otherwise, we can reuse the precomputed value
+            consistent = False
+
+        if not consistent:
+            self.cached_chunk_size = self._determine_favorable_chunk_size(
+                representative_fn,
+                args,
+                min_chunk_size,
+            )
+            self.cached_arg_data = arg_data
+
+        assert self.cached_chunk_size is not None
+
+        return self.cached_chunk_size

LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/esm/openfold_utils/loss.py

@@ -0,0 +1,104 @@
+# Copyright 2021 AlQuraishi Laboratory
+# Copyright 2021 DeepMind Technologies Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import torch
+
+
+def _calculate_bin_centers(boundaries: torch.Tensor) -> torch.Tensor:
+    step = boundaries[1] - boundaries[0]
+    bin_centers = boundaries + step / 2
+    bin_centers = torch.cat([bin_centers, (bin_centers[-1] + step).unsqueeze(-1)], dim=0)
+    return bin_centers
+
+
+def _calculate_expected_aligned_error(
+    alignment_confidence_breaks: torch.Tensor,
+    aligned_distance_error_probs: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    bin_centers = _calculate_bin_centers(alignment_confidence_breaks)
+    return (
+        torch.sum(aligned_distance_error_probs * bin_centers, dim=-1),
+        bin_centers[-1],
+    )
+
+
+def compute_predicted_aligned_error(
+    logits: torch.Tensor,
+    max_bin: int = 31,
+    no_bins: int = 64,
+    **kwargs,
+) -> dict[str, torch.Tensor]:
+    """Computes aligned confidence metrics from logits.
+
+    Args:
+      logits: [*, num_res, num_res, num_bins] the logits output from
+        PredictedAlignedErrorHead.
+      max_bin: Maximum bin value
+      no_bins: Number of bins
+    Returns:
+      aligned_confidence_probs: [*, num_res, num_res, num_bins] the predicted
+        aligned error probabilities over bins for each residue pair.
+      predicted_aligned_error: [*, num_res, num_res] the expected aligned distance
+        error for each pair of residues.
+      max_predicted_aligned_error: [*] the maximum predicted error possible.
+    """
+    boundaries = torch.linspace(0, max_bin, steps=(no_bins - 1), device=logits.device)
+
+    aligned_confidence_probs = torch.nn.functional.softmax(logits, dim=-1)
+    predicted_aligned_error, max_predicted_aligned_error = _calculate_expected_aligned_error(
+        alignment_confidence_breaks=boundaries,
+        aligned_distance_error_probs=aligned_confidence_probs,
+    )
+
+    return {
+        "aligned_confidence_probs": aligned_confidence_probs,
+        "predicted_aligned_error": predicted_aligned_error,
+        "max_predicted_aligned_error": max_predicted_aligned_error,
+    }
+
+
+def compute_tm(
+    logits: torch.Tensor,
+    residue_weights: torch.Tensor | None = None,
+    max_bin: int = 31,
+    no_bins: int = 64,
+    eps: float = 1e-8,
+    **kwargs,
+) -> torch.Tensor:
+    if residue_weights is None:
+        residue_weights = logits.new_ones(logits.shape[-2])
+
+    boundaries = torch.linspace(0, max_bin, steps=(no_bins - 1), device=logits.device)
+
+    bin_centers = _calculate_bin_centers(boundaries)
+    torch.sum(residue_weights)
+    n = logits.shape[-2]
+    clipped_n = max(n, 19)
+
+    d0 = 1.24 * (clipped_n - 15) ** (1.0 / 3) - 1.8
+
+    probs = torch.nn.functional.softmax(logits, dim=-1)
+
+    tm_per_bin = 1.0 / (1 + (bin_centers**2) / (d0**2))
+    predicted_tm_term = torch.sum(probs * tm_per_bin, dim=-1)
+
+    normed_residue_mask = residue_weights / (eps + residue_weights.sum())
+    per_alignment = torch.sum(predicted_tm_term * normed_residue_mask, dim=-1)
+
+    weighted = per_alignment * residue_weights
+
+    argmax = (weighted == torch.max(weighted)).nonzero()[0]
+    return per_alignment[tuple(argmax)]

LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/esm/openfold_utils/protein.py

@@ -0,0 +1,330 @@
+# Copyright 2021 AlQuraishi Laboratory
+# Copyright 2021 DeepMind Technologies Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Protein data type."""
+
+import dataclasses
+import re
+import string
+from collections.abc import Iterator, Mapping, Sequence
+from typing import Any
+
+import numpy as np
+
+from . import residue_constants
+
+
+FeatureDict = Mapping[str, np.ndarray]
+ModelOutput = Mapping[str, Any]  # Is a nested dict.
+PICO_TO_ANGSTROM = 0.01
+
+
+@dataclasses.dataclass(frozen=True)
+class Protein:
+    """Protein structure representation."""
+
+    # Cartesian coordinates of atoms in angstroms. The atom types correspond to
+    # residue_constants.atom_types, i.e. the first three are N, CA, CB.
+    atom_positions: np.ndarray  # [num_res, num_atom_type, 3]
+
+    # Amino-acid type for each residue represented as an integer between 0 and
+    # 20, where 20 is 'X'.
+    aatype: np.ndarray  # [num_res]
+
+    # Binary float mask to indicate presence of a particular atom. 1.0 if an atom
+    # is present and 0.0 if not. This should be used for loss masking.
+    atom_mask: np.ndarray  # [num_res, num_atom_type]
+
+    # Residue index as used in PDB. It is not necessarily continuous or 0-indexed.
+    residue_index: np.ndarray  # [num_res]
+
+    # B-factors, or temperature factors, of each residue (in sq. angstroms units),
+    # representing the displacement of the residue from its ground truth mean
+    # value.
+    b_factors: np.ndarray  # [num_res, num_atom_type]
+
+    # Chain indices for multi-chain predictions
+    chain_index: np.ndarray | None = None
+
+    # Optional remark about the protein. Included as a comment in output PDB
+    # files
+    remark: str | None = None
+
+    # Templates used to generate this protein (prediction-only)
+    parents: Sequence[str] | None = None
+
+    # Chain corresponding to each parent
+    parents_chain_index: Sequence[int] | None = None
+
+
+def from_proteinnet_string(proteinnet_str: str) -> Protein:
+    tag_re = r"(\[[A-Z]+\]\n)"
+    tags: list[str] = [tag.strip() for tag in re.split(tag_re, proteinnet_str) if len(tag) > 0]
+    groups: Iterator[tuple[str, list[str]]] = zip(tags[0::2], [l.split("\n") for l in tags[1::2]])
+
+    atoms: list[str] = ["N", "CA", "C"]
+    aatype = None
+    atom_positions = None
+    atom_mask = None
+    for g in groups:
+        if g[0] == "[PRIMARY]":
+            seq = g[1][0].strip()
+            # Replace unknown residues with "X" (strings are immutable, so convert to list first)
+            seq = [char if char in residue_constants.restypes else "X" for char in seq]
+            aatype = np.array(
+                [residue_constants.restype_order.get(res_symbol, residue_constants.restype_num) for res_symbol in seq]
+            )
+        elif g[0] == "[TERTIARY]":
+            tertiary: list[list[float]] = []
+            for axis in range(3):
+                tertiary.append(list(map(float, g[1][axis].split())))
+            tertiary_np = np.array(tertiary)
+            atom_positions = np.zeros((len(tertiary[0]) // 3, residue_constants.atom_type_num, 3)).astype(np.float32)
+            for i, atom in enumerate(atoms):
+                atom_positions[:, residue_constants.atom_order[atom], :] = np.transpose(tertiary_np[:, i::3])
+            atom_positions *= PICO_TO_ANGSTROM
+        elif g[0] == "[MASK]":
+            mask = np.array(list(map({"-": 0, "+": 1}.get, g[1][0].strip())))
+            atom_mask = np.zeros(
+                (
+                    len(mask),
+                    residue_constants.atom_type_num,
+                )
+            ).astype(np.float32)
+            for i, atom in enumerate(atoms):
+                atom_mask[:, residue_constants.atom_order[atom]] = 1
+            atom_mask *= mask[..., None]
+
+    assert aatype is not None
+
+    return Protein(
+        atom_positions=atom_positions,
+        atom_mask=atom_mask,
+        aatype=aatype,
+        residue_index=np.arange(len(aatype)),
+        b_factors=None,
+    )
+
+
+def get_pdb_headers(prot: Protein, chain_id: int = 0) -> list[str]:
+    pdb_headers: list[str] = []
+
+    remark = prot.remark
+    if remark is not None:
+        pdb_headers.append(f"REMARK {remark}")
+
+    parents = prot.parents
+    parents_chain_index = prot.parents_chain_index
+    if parents is not None and parents_chain_index is not None:
+        parents = [p for i, p in zip(parents_chain_index, parents) if i == chain_id]
+
+    if parents is None or len(parents) == 0:
+        parents = ["N/A"]
+
+    pdb_headers.append(f"PARENT {' '.join(parents)}")
+
+    return pdb_headers
+
+
+def add_pdb_headers(prot: Protein, pdb_str: str) -> str:
+    """Add pdb headers to an existing PDB string. Useful during multi-chain
+    recycling
+    """
+    out_pdb_lines: list[str] = []
+    lines = pdb_str.split("\n")
+
+    remark = prot.remark
+    if remark is not None:
+        out_pdb_lines.append(f"REMARK {remark}")
+
+    parents_per_chain: list[list[str]]
+    if prot.parents is not None and len(prot.parents) > 0:
+        parents_per_chain = []
+        if prot.parents_chain_index is not None:
+            parent_dict: dict[str, list[str]] = {}
+            for p, i in zip(prot.parents, prot.parents_chain_index):
+                parent_dict.setdefault(str(i), [])
+                parent_dict[str(i)].append(p)
+
+            max_idx = max(int(chain_idx) for chain_idx in parent_dict)
+            for i in range(max_idx + 1):
+                chain_parents = parent_dict.get(str(i), ["N/A"])
+                parents_per_chain.append(chain_parents)
+        else:
+            parents_per_chain.append(list(prot.parents))
+    else:
+        parents_per_chain = [["N/A"]]
+
+    def make_parent_line(p: Sequence[str]) -> str:
+        return f"PARENT {' '.join(p)}"
+
+    out_pdb_lines.append(make_parent_line(parents_per_chain[0]))
+
+    chain_counter = 0
+    for i, l in enumerate(lines):
+        if "PARENT" not in l and "REMARK" not in l:
+            out_pdb_lines.append(l)
+        if "TER" in l and "END" not in lines[i + 1]:
+            chain_counter += 1
+            if not chain_counter >= len(parents_per_chain):
+                chain_parents = parents_per_chain[chain_counter]
+            else:
+                chain_parents = ["N/A"]
+
+            out_pdb_lines.append(make_parent_line(chain_parents))
+
+    return "\n".join(out_pdb_lines)
+
+
+def to_pdb(prot: Protein) -> str:
+    """Converts a `Protein` instance to a PDB string.
+
+    Args:
+      prot: The protein to convert to PDB.
+
+    Returns:
+      PDB string.
+    """
+    restypes = residue_constants.restypes + ["X"]
+
+    def res_1to3(r: int) -> str:
+        return residue_constants.restype_1to3.get(restypes[r], "UNK")
+
+    atom_types = residue_constants.atom_types
+
+    pdb_lines: list[str] = []
+
+    atom_mask = prot.atom_mask
+    aatype = prot.aatype
+    atom_positions = prot.atom_positions
+    residue_index = prot.residue_index.astype(np.int32)
+    b_factors = prot.b_factors
+    chain_index = prot.chain_index
+
+    if np.any(aatype > residue_constants.restype_num):
+        raise ValueError("Invalid aatypes.")
+
+    headers = get_pdb_headers(prot)
+    if len(headers) > 0:
+        pdb_lines.extend(headers)
+
+    n = aatype.shape[0]
+    atom_index = 1
+    prev_chain_index = 0
+    chain_tags = string.ascii_uppercase
+    chain_tag = None
+    # Add all atom sites.
+    for i in range(n):
+        res_name_3 = res_1to3(aatype[i])
+        for atom_name, pos, mask, b_factor in zip(atom_types, atom_positions[i], atom_mask[i], b_factors[i]):
+            if mask < 0.5:
+                continue
+
+            record_type = "ATOM"
+            name = atom_name if len(atom_name) == 4 else f" {atom_name}"
+            alt_loc = ""
+            insertion_code = ""
+            occupancy = 1.00
+            element = atom_name[0]  # Protein supports only C, N, O, S, this works.
+            charge = ""
+
+            chain_tag = "A"
+            if chain_index is not None:
+                chain_tag = chain_tags[chain_index[i]]
+
+            # PDB is a columnar format, every space matters here!
+            atom_line = (
+                f"{record_type:<6}{atom_index:>5} {name:<4}{alt_loc:>1}"
+                f"{res_name_3:>3} {chain_tag:>1}"
+                f"{residue_index[i]:>4}{insertion_code:>1}   "
+                f"{pos[0]:>8.3f}{pos[1]:>8.3f}{pos[2]:>8.3f}"
+                f"{occupancy:>6.2f}{b_factor:>6.2f}          "
+                f"{element:>2}{charge:>2}"
+            )
+            pdb_lines.append(atom_line)
+            atom_index += 1
+
+        should_terminate = i == n - 1
+        if chain_index is not None:
+            if i != n - 1 and chain_index[i + 1] != prev_chain_index:
+                should_terminate = True
+                prev_chain_index = chain_index[i + 1]
+
+        if should_terminate:
+            # Close the chain.
+            chain_end = "TER"
+            chain_termination_line = (
+                f"{chain_end:<6}{atom_index:>5}      {res_1to3(aatype[i]):>3} {chain_tag:>1}{residue_index[i]:>4}"
+            )
+            pdb_lines.append(chain_termination_line)
+            atom_index += 1
+
+            if i != n - 1:
+                # "prev" is a misnomer here. This happens at the beginning of
+                # each new chain.
+                pdb_lines.extend(get_pdb_headers(prot, prev_chain_index))
+
+    pdb_lines.append("END")
+    pdb_lines.append("")
+    return "\n".join(pdb_lines)
+
+
+def ideal_atom_mask(prot: Protein) -> np.ndarray:
+    """Computes an ideal atom mask.
+
+    `Protein.atom_mask` typically is defined according to the atoms that are reported in the PDB. This function
+    computes a mask according to heavy atoms that should be present in the given sequence of amino acids.
+
+    Args:
+      prot: `Protein` whose fields are `numpy.ndarray` objects.
+
+    Returns:
+      An ideal atom mask.
+    """
+    return residue_constants.STANDARD_ATOM_MASK[prot.aatype]
+
+
+def from_prediction(
+    features: FeatureDict,
+    result: ModelOutput,
+    b_factors: np.ndarray | None = None,
+    chain_index: np.ndarray | None = None,
+    remark: str | None = None,
+    parents: Sequence[str] | None = None,
+    parents_chain_index: Sequence[int] | None = None,
+) -> Protein:
+    """Assembles a protein from a prediction.
+
+    Args:
+      features: Dictionary holding model inputs.
+      result: Dictionary holding model outputs.
+      b_factors: (Optional) B-factors to use for the protein.
+      chain_index: (Optional) Chain indices for multi-chain predictions
+      remark: (Optional) Remark about the prediction
+      parents: (Optional) List of template names
+    Returns:
+      A protein instance.
+    """
+    return Protein(
+        aatype=features["aatype"],
+        atom_positions=result["final_atom_positions"],
+        atom_mask=result["final_atom_mask"],
+        residue_index=features["residue_index"] + 1,
+        b_factors=b_factors if b_factors is not None else np.zeros_like(result["final_atom_mask"]),
+        chain_index=chain_index,
+        remark=remark,
+        parents=parents,
+        parents_chain_index=parents_chain_index,
+    )

LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/esm/openfold_utils/residue_constants.py

@@ -0,0 +1,979 @@
+# Copyright 2021 AlQuraishi Laboratory
+# Copyright 2021 DeepMind Technologies Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Constants used in AlphaFold."""
+
+import collections
+import copy
+import functools
+from collections.abc import Mapping, Sequence
+from importlib import resources
+
+import numpy as np
+
+
+# Internal import (35fd).
+
+
+# Distance from one CA to next CA [trans configuration: omega = 180].
+ca_ca = 3.80209737096
+
+# Format: The list for each AA type contains chi1, chi2, chi3, chi4 in
+# this order (or a relevant subset from chi1 onwards). ALA and GLY don't have
+# chi angles so their chi angle lists are empty.
+chi_angles_atoms: dict[str, list[list[str]]] = {
+    "ALA": [],
+    # Chi5 in arginine is always 0 +- 5 degrees, so ignore it.
+    "ARG": [["N", "CA", "CB", "CG"], ["CA", "CB", "CG", "CD"], ["CB", "CG", "CD", "NE"], ["CG", "CD", "NE", "CZ"]],
+    "ASN": [["N", "CA", "CB", "CG"], ["CA", "CB", "CG", "OD1"]],
+    "ASP": [["N", "CA", "CB", "CG"], ["CA", "CB", "CG", "OD1"]],
+    "CYS": [["N", "CA", "CB", "SG"]],
+    "GLN": [["N", "CA", "CB", "CG"], ["CA", "CB", "CG", "CD"], ["CB", "CG", "CD", "OE1"]],
+    "GLU": [["N", "CA", "CB", "CG"], ["CA", "CB", "CG", "CD"], ["CB", "CG", "CD", "OE1"]],
+    "GLY": [],
+    "HIS": [["N", "CA", "CB", "CG"], ["CA", "CB", "CG", "ND1"]],
+    "ILE": [["N", "CA", "CB", "CG1"], ["CA", "CB", "CG1", "CD1"]],
+    "LEU": [["N", "CA", "CB", "CG"], ["CA", "CB", "CG", "CD1"]],
+    "LYS": [["N", "CA", "CB", "CG"], ["CA", "CB", "CG", "CD"], ["CB", "CG", "CD", "CE"], ["CG", "CD", "CE", "NZ"]],
+    "MET": [["N", "CA", "CB", "CG"], ["CA", "CB", "CG", "SD"], ["CB", "CG", "SD", "CE"]],
+    "PHE": [["N", "CA", "CB", "CG"], ["CA", "CB", "CG", "CD1"]],
+    "PRO": [["N", "CA", "CB", "CG"], ["CA", "CB", "CG", "CD"]],
+    "SER": [["N", "CA", "CB", "OG"]],
+    "THR": [["N", "CA", "CB", "OG1"]],
+    "TRP": [["N", "CA", "CB", "CG"], ["CA", "CB", "CG", "CD1"]],
+    "TYR": [["N", "CA", "CB", "CG"], ["CA", "CB", "CG", "CD1"]],
+    "VAL": [["N", "CA", "CB", "CG1"]],
+}
+
+# If chi angles given in fixed-length array, this matrix determines how to mask
+# them for each AA type. The order is as per restype_order (see below).
+chi_angles_mask: list[list[float]] = [
+    [0.0, 0.0, 0.0, 0.0],  # ALA
+    [1.0, 1.0, 1.0, 1.0],  # ARG
+    [1.0, 1.0, 0.0, 0.0],  # ASN
+    [1.0, 1.0, 0.0, 0.0],  # ASP
+    [1.0, 0.0, 0.0, 0.0],  # CYS
+    [1.0, 1.0, 1.0, 0.0],  # GLN
+    [1.0, 1.0, 1.0, 0.0],  # GLU
+    [0.0, 0.0, 0.0, 0.0],  # GLY
+    [1.0, 1.0, 0.0, 0.0],  # HIS
+    [1.0, 1.0, 0.0, 0.0],  # ILE
+    [1.0, 1.0, 0.0, 0.0],  # LEU
+    [1.0, 1.0, 1.0, 1.0],  # LYS
+    [1.0, 1.0, 1.0, 0.0],  # MET
+    [1.0, 1.0, 0.0, 0.0],  # PHE
+    [1.0, 1.0, 0.0, 0.0],  # PRO
+    [1.0, 0.0, 0.0, 0.0],  # SER
+    [1.0, 0.0, 0.0, 0.0],  # THR
+    [1.0, 1.0, 0.0, 0.0],  # TRP
+    [1.0, 1.0, 0.0, 0.0],  # TYR
+    [1.0, 0.0, 0.0, 0.0],  # VAL
+]
+
+# The following chi angles are pi periodic: they can be rotated by a multiple
+# of pi without affecting the structure.
+chi_pi_periodic: list[list[float]] = [
+    [0.0, 0.0, 0.0, 0.0],  # ALA
+    [0.0, 0.0, 0.0, 0.0],  # ARG
+    [0.0, 0.0, 0.0, 0.0],  # ASN
+    [0.0, 1.0, 0.0, 0.0],  # ASP
+    [0.0, 0.0, 0.0, 0.0],  # CYS
+    [0.0, 0.0, 0.0, 0.0],  # GLN
+    [0.0, 0.0, 1.0, 0.0],  # GLU
+    [0.0, 0.0, 0.0, 0.0],  # GLY
+    [0.0, 0.0, 0.0, 0.0],  # HIS
+    [0.0, 0.0, 0.0, 0.0],  # ILE
+    [0.0, 0.0, 0.0, 0.0],  # LEU
+    [0.0, 0.0, 0.0, 0.0],  # LYS
+    [0.0, 0.0, 0.0, 0.0],  # MET
+    [0.0, 1.0, 0.0, 0.0],  # PHE
+    [0.0, 0.0, 0.0, 0.0],  # PRO
+    [0.0, 0.0, 0.0, 0.0],  # SER
+    [0.0, 0.0, 0.0, 0.0],  # THR
+    [0.0, 0.0, 0.0, 0.0],  # TRP
+    [0.0, 1.0, 0.0, 0.0],  # TYR
+    [0.0, 0.0, 0.0, 0.0],  # VAL
+    [0.0, 0.0, 0.0, 0.0],  # UNK
+]
+
+# Atoms positions relative to the 8 rigid groups, defined by the pre-omega, phi,
+# psi and chi angles:
+# 0: 'backbone group',
+# 1: 'pre-omega-group', (empty)
+# 2: 'phi-group', (currently empty, because it defines only hydrogens)
+# 3: 'psi-group',
+# 4,5,6,7: 'chi1,2,3,4-group'
+# The atom positions are relative to the axis-end-atom of the corresponding
+# rotation axis. The x-axis is in direction of the rotation axis, and the y-axis
+# is defined such that the dihedral-angle-definiting atom (the last entry in
+# chi_angles_atoms above) is in the xy-plane (with a positive y-coordinate).
+# format: [atomname, group_idx, rel_position]
+rigid_group_atom_positions: dict[str, list[tuple[str, int, tuple[float, float, float]]]] = {
+    "ALA": [
+        ("N", 0, (-0.525, 1.363, 0.000)),
+        ("CA", 0, (0.000, 0.000, 0.000)),
+        ("C", 0, (1.526, -0.000, -0.000)),
+        ("CB", 0, (-0.529, -0.774, -1.205)),
+        ("O", 3, (0.627, 1.062, 0.000)),
+    ],
+    "ARG": [
+        ("N", 0, (-0.524, 1.362, -0.000)),
+        ("CA", 0, (0.000, 0.000, 0.000)),
+        ("C", 0, (1.525, -0.000, -0.000)),
+        ("CB", 0, (-0.524, -0.778, -1.209)),
+        ("O", 3, (0.626, 1.062, 0.000)),
+        ("CG", 4, (0.616, 1.390, -0.000)),
+        ("CD", 5, (0.564, 1.414, 0.000)),
+        ("NE", 6, (0.539, 1.357, -0.000)),
+        ("NH1", 7, (0.206, 2.301, 0.000)),
+        ("NH2", 7, (2.078, 0.978, -0.000)),
+        ("CZ", 7, (0.758, 1.093, -0.000)),
+    ],
+    "ASN": [
+        ("N", 0, (-0.536, 1.357, 0.000)),
+        ("CA", 0, (0.000, 0.000, 0.000)),
+        ("C", 0, (1.526, -0.000, -0.000)),
+        ("CB", 0, (-0.531, -0.787, -1.200)),
+        ("O", 3, (0.625, 1.062, 0.000)),
+        ("CG", 4, (0.584, 1.399, 0.000)),
+        ("ND2", 5, (0.593, -1.188, 0.001)),
+        ("OD1", 5, (0.633, 1.059, 0.000)),
+    ],
+    "ASP": [
+        ("N", 0, (-0.525, 1.362, -0.000)),
+        ("CA", 0, (0.000, 0.000, 0.000)),
+        ("C", 0, (1.527, 0.000, -0.000)),
+        ("CB", 0, (-0.526, -0.778, -1.208)),
+        ("O", 3, (0.626, 1.062, -0.000)),
+        ("CG", 4, (0.593, 1.398, -0.000)),
+        ("OD1", 5, (0.610, 1.091, 0.000)),
+        ("OD2", 5, (0.592, -1.101, -0.003)),
+    ],
+    "CYS": [
+        ("N", 0, (-0.522, 1.362, -0.000)),
+        ("CA", 0, (0.000, 0.000, 0.000)),
+        ("C", 0, (1.524, 0.000, 0.000)),
+        ("CB", 0, (-0.519, -0.773, -1.212)),
+        ("O", 3, (0.625, 1.062, -0.000)),
+        ("SG", 4, (0.728, 1.653, 0.000)),
+    ],
+    "GLN": [
+        ("N", 0, (-0.526, 1.361, -0.000)),
+        ("CA", 0, (0.000, 0.000, 0.000)),
+        ("C", 0, (1.526, 0.000, 0.000)),
+        ("CB", 0, (-0.525, -0.779, -1.207)),
+        ("O", 3, (0.626, 1.062, -0.000)),
+        ("CG", 4, (0.615, 1.393, 0.000)),
+        ("CD", 5, (0.587, 1.399, -0.000)),
+        ("NE2", 6, (0.593, -1.189, -0.001)),
+        ("OE1", 6, (0.634, 1.060, 0.000)),
+    ],
+    "GLU": [
+        ("N", 0, (-0.528, 1.361, 0.000)),
+        ("CA", 0, (0.000, 0.000, 0.000)),
+        ("C", 0, (1.526, -0.000, -0.000)),
+        ("CB", 0, (-0.526, -0.781, -1.207)),
+        ("O", 3, (0.626, 1.062, 0.000)),
+        ("CG", 4, (0.615, 1.392, 0.000)),
+        ("CD", 5, (0.600, 1.397, 0.000)),
+        ("OE1", 6, (0.607, 1.095, -0.000)),
+        ("OE2", 6, (0.589, -1.104, -0.001)),
+    ],
+    "GLY": [
+        ("N", 0, (-0.572, 1.337, 0.000)),
+        ("CA", 0, (0.000, 0.000, 0.000)),
+        ("C", 0, (1.517, -0.000, -0.000)),
+        ("O", 3, (0.626, 1.062, -0.000)),
+    ],
+    "HIS": [
+        ("N", 0, (-0.527, 1.360, 0.000)),
+        ("CA", 0, (0.000, 0.000, 0.000)),
+        ("C", 0, (1.525, 0.000, 0.000)),
+        ("CB", 0, (-0.525, -0.778, -1.208)),
+        ("O", 3, (0.625, 1.063, 0.000)),
+        ("CG", 4, (0.600, 1.370, -0.000)),
+        ("CD2", 5, (0.889, -1.021, 0.003)),
+        ("ND1", 5, (0.744, 1.160, -0.000)),
+        ("CE1", 5, (2.030, 0.851, 0.002)),
+        ("NE2", 5, (2.145, -0.466, 0.004)),
+    ],
+    "ILE": [
+        ("N", 0, (-0.493, 1.373, -0.000)),
+        ("CA", 0, (0.000, 0.000, 0.000)),
+        ("C", 0, (1.527, -0.000, -0.000)),
+        ("CB", 0, (-0.536, -0.793, -1.213)),
+        ("O", 3, (0.627, 1.062, -0.000)),
+        ("CG1", 4, (0.534, 1.437, -0.000)),
+        ("CG2", 4, (0.540, -0.785, -1.199)),
+        ("CD1", 5, (0.619, 1.391, 0.000)),
+    ],
+    "LEU": [
+        ("N", 0, (-0.520, 1.363, 0.000)),
+        ("CA", 0, (0.000, 0.000, 0.000)),
+        ("C", 0, (1.525, -0.000, -0.000)),
+        ("CB", 0, (-0.522, -0.773, -1.214)),
+        ("O", 3, (0.625, 1.063, -0.000)),
+        ("CG", 4, (0.678, 1.371, 0.000)),
+        ("CD1", 5, (0.530, 1.430, -0.000)),
+        ("CD2", 5, (0.535, -0.774, 1.200)),
+    ],
+    "LYS": [
+        ("N", 0, (-0.526, 1.362, -0.000)),
+        ("CA", 0, (0.000, 0.000, 0.000)),
+        ("C", 0, (1.526, 0.000, 0.000)),
+        ("CB", 0, (-0.524, -0.778, -1.208)),
+        ("O", 3, (0.626, 1.062, -0.000)),
+        ("CG", 4, (0.619, 1.390, 0.000)),
+        ("CD", 5, (0.559, 1.417, 0.000)),
+        ("CE", 6, (0.560, 1.416, 0.000)),
+        ("NZ", 7, (0.554, 1.387, 0.000)),
+    ],
+    "MET": [
+        ("N", 0, (-0.521, 1.364, -0.000)),
+        ("CA", 0, (0.000, 0.000, 0.000)),
+        ("C", 0, (1.525, 0.000, 0.000)),
+        ("CB", 0, (-0.523, -0.776, -1.210)),
+        ("O", 3, (0.625, 1.062, -0.000)),
+        ("CG", 4, (0.613, 1.391, -0.000)),
+        ("SD", 5, (0.703, 1.695, 0.000)),
+        ("CE", 6, (0.320, 1.786, -0.000)),
+    ],
+    "PHE": [
+        ("N", 0, (-0.518, 1.363, 0.000)),
+        ("CA", 0, (0.000, 0.000, 0.000)),
+        ("C", 0, (1.524, 0.000, -0.000)),
+        ("CB", 0, (-0.525, -0.776, -1.212)),
+        ("O", 3, (0.626, 1.062, -0.000)),
+        ("CG", 4, (0.607, 1.377, 0.000)),
+        ("CD1", 5, (0.709, 1.195, -0.000)),
+        ("CD2", 5, (0.706, -1.196, 0.000)),
+        ("CE1", 5, (2.102, 1.198, -0.000)),
+        ("CE2", 5, (2.098, -1.201, -0.000)),
+        ("CZ", 5, (2.794, -0.003, -0.001)),
+    ],
+    "PRO": [
+        ("N", 0, (-0.566, 1.351, -0.000)),
+        ("CA", 0, (0.000, 0.000, 0.000)),
+        ("C", 0, (1.527, -0.000, 0.000)),
+        ("CB", 0, (-0.546, -0.611, -1.293)),
+        ("O", 3, (0.621, 1.066, 0.000)),
+        ("CG", 4, (0.382, 1.445, 0.0)),
+        # ('CD', 5, (0.427, 1.440, 0.0)),
+        ("CD", 5, (0.477, 1.424, 0.0)),  # manually made angle 2 degrees larger
+    ],
+    "SER": [
+        ("N", 0, (-0.529, 1.360, -0.000)),
+        ("CA", 0, (0.000, 0.000, 0.000)),
+        ("C", 0, (1.525, -0.000, -0.000)),
+        ("CB", 0, (-0.518, -0.777, -1.211)),
+        ("O", 3, (0.626, 1.062, -0.000)),
+        ("OG", 4, (0.503, 1.325, 0.000)),
+    ],
+    "THR": [
+        ("N", 0, (-0.517, 1.364, 0.000)),
+        ("CA", 0, (0.000, 0.000, 0.000)),
+        ("C", 0, (1.526, 0.000, -0.000)),
+        ("CB", 0, (-0.516, -0.793, -1.215)),
+        ("O", 3, (0.626, 1.062, 0.000)),
+        ("CG2", 4, (0.550, -0.718, -1.228)),
+        ("OG1", 4, (0.472, 1.353, 0.000)),
+    ],
+    "TRP": [
+        ("N", 0, (-0.521, 1.363, 0.000)),
+        ("CA", 0, (0.000, 0.000, 0.000)),
+        ("C", 0, (1.525, -0.000, 0.000)),
+        ("CB", 0, (-0.523, -0.776, -1.212)),
+        ("O", 3, (0.627, 1.062, 0.000)),
+        ("CG", 4, (0.609, 1.370, -0.000)),
+        ("CD1", 5, (0.824, 1.091, 0.000)),
+        ("CD2", 5, (0.854, -1.148, -0.005)),
+        ("CE2", 5, (2.186, -0.678, -0.007)),
+        ("CE3", 5, (0.622, -2.530, -0.007)),
+        ("NE1", 5, (2.140, 0.690, -0.004)),
+        ("CH2", 5, (3.028, -2.890, -0.013)),
+        ("CZ2", 5, (3.283, -1.543, -0.011)),
+        ("CZ3", 5, (1.715, -3.389, -0.011)),
+    ],
+    "TYR": [
+        ("N", 0, (-0.522, 1.362, 0.000)),
+        ("CA", 0, (0.000, 0.000, 0.000)),
+        ("C", 0, (1.524, -0.000, -0.000)),
+        ("CB", 0, (-0.522, -0.776, -1.213)),
+        ("O", 3, (0.627, 1.062, -0.000)),
+        ("CG", 4, (0.607, 1.382, -0.000)),
+        ("CD1", 5, (0.716, 1.195, -0.000)),
+        ("CD2", 5, (0.713, -1.194, -0.001)),
+        ("CE1", 5, (2.107, 1.200, -0.002)),
+        ("CE2", 5, (2.104, -1.201, -0.003)),
+        ("OH", 5, (4.168, -0.002, -0.005)),
+        ("CZ", 5, (2.791, -0.001, -0.003)),
+    ],
+    "VAL": [
+        ("N", 0, (-0.494, 1.373, -0.000)),
+        ("CA", 0, (0.000, 0.000, 0.000)),
+        ("C", 0, (1.527, -0.000, -0.000)),
+        ("CB", 0, (-0.533, -0.795, -1.213)),
+        ("O", 3, (0.627, 1.062, -0.000)),
+        ("CG1", 4, (0.540, 1.429, -0.000)),
+        ("CG2", 4, (0.533, -0.776, 1.203)),
+    ],
+}
+
+# A list of atoms (excluding hydrogen) for each AA type. PDB naming convention.
+residue_atoms: dict[str, list[str]] = {
+    "ALA": ["C", "CA", "CB", "N", "O"],
+    "ARG": ["C", "CA", "CB", "CG", "CD", "CZ", "N", "NE", "O", "NH1", "NH2"],
+    "ASP": ["C", "CA", "CB", "CG", "N", "O", "OD1", "OD2"],
+    "ASN": ["C", "CA", "CB", "CG", "N", "ND2", "O", "OD1"],
+    "CYS": ["C", "CA", "CB", "N", "O", "SG"],
+    "GLU": ["C", "CA", "CB", "CG", "CD", "N", "O", "OE1", "OE2"],
+    "GLN": ["C", "CA", "CB", "CG", "CD", "N", "NE2", "O", "OE1"],
+    "GLY": ["C", "CA", "N", "O"],
+    "HIS": ["C", "CA", "CB", "CG", "CD2", "CE1", "N", "ND1", "NE2", "O"],
+    "ILE": ["C", "CA", "CB", "CG1", "CG2", "CD1", "N", "O"],
+    "LEU": ["C", "CA", "CB", "CG", "CD1", "CD2", "N", "O"],
+    "LYS": ["C", "CA", "CB", "CG", "CD", "CE", "N", "NZ", "O"],
+    "MET": ["C", "CA", "CB", "CG", "CE", "N", "O", "SD"],
+    "PHE": ["C", "CA", "CB", "CG", "CD1", "CD2", "CE1", "CE2", "CZ", "N", "O"],
+    "PRO": ["C", "CA", "CB", "CG", "CD", "N", "O"],
+    "SER": ["C", "CA", "CB", "N", "O", "OG"],
+    "THR": ["C", "CA", "CB", "CG2", "N", "O", "OG1"],
+    "TRP": ["C", "CA", "CB", "CG", "CD1", "CD2", "CE2", "CE3", "CZ2", "CZ3", "CH2", "N", "NE1", "O"],
+    "TYR": ["C", "CA", "CB", "CG", "CD1", "CD2", "CE1", "CE2", "CZ", "N", "O", "OH"],
+    "VAL": ["C", "CA", "CB", "CG1", "CG2", "N", "O"],
+}
+
+# Naming swaps for ambiguous atom names.
+# Due to symmetries in the amino acids the naming of atoms is ambiguous in
+# 4 of the 20 amino acids.
+# (The LDDT paper lists 7 amino acids as ambiguous, but the naming ambiguities
+# in LEU, VAL and ARG can be resolved by using the 3d constellations of
+# the 'ambiguous' atoms and their neighbours)
+# TODO: ^ interpret this
+residue_atom_renaming_swaps: dict[str, dict[str, str]] = {
+    "ASP": {"OD1": "OD2"},
+    "GLU": {"OE1": "OE2"},
+    "PHE": {"CD1": "CD2", "CE1": "CE2"},
+    "TYR": {"CD1": "CD2", "CE1": "CE2"},
+}
+
+# Van der Waals radii [Angstroem] of the atoms (from Wikipedia)
+van_der_waals_radius: dict[str, float] = {
+    "C": 1.7,
+    "N": 1.55,
+    "O": 1.52,
+    "S": 1.8,
+}
+
+Bond = collections.namedtuple("Bond", ["atom1_name", "atom2_name", "length", "stddev"])
+BondAngle = collections.namedtuple(
+    "BondAngle",
+    ["atom1_name", "atom2_name", "atom3name", "angle_rad", "stddev"],
+)
+
+
+def map_structure_with_atom_order(in_list: list, first_call: bool = True) -> list:
+    # Maps strings in a nested list structure to their corresponding index in atom_order
+    if first_call:
+        in_list = copy.deepcopy(in_list)
+    for i in range(len(in_list)):
+        if isinstance(in_list[i], list):
+            in_list[i] = map_structure_with_atom_order(in_list[i], first_call=False)
+        elif isinstance(in_list[i], str):
+            in_list[i] = atom_order[in_list[i]]
+        else:
+            raise TypeError("Unexpected type when mapping nested lists!")
+    return in_list
+
+
+@functools.cache
+def load_stereo_chemical_props() -> tuple[
+    Mapping[str, list[Bond]],
+    Mapping[str, list[Bond]],
+    Mapping[str, list[BondAngle]],
+]:
+    """Load stereo_chemical_props.txt into a nice structure.
+
+    Load literature values for bond lengths and bond angles and translate bond angles into the length of the opposite
+    edge of the triangle ("residue_virtual_bonds").
+
+    Returns:
+      residue_bonds: dict that maps resname --> list of Bond tuples residue_virtual_bonds: dict that maps resname -->
+      list of Bond tuples residue_bond_angles: dict that maps resname --> list of BondAngle tuples
+    """
+    # TODO: this file should be downloaded in a setup script
+    stereo_chemical_props = resources.read_text("openfold.resources", "stereo_chemical_props.txt")
+
+    lines_iter = iter(stereo_chemical_props.splitlines())
+    # Load bond lengths.
+    residue_bonds: dict[str, list[Bond]] = {}
+    next(lines_iter)  # Skip header line.
+    for line in lines_iter:
+        if line.strip() == "-":
+            break
+        bond, resname, bond_length, stddev = line.split()
+        atom1, atom2 = bond.split("-")
+        if resname not in residue_bonds:
+            residue_bonds[resname] = []
+        residue_bonds[resname].append(Bond(atom1, atom2, float(bond_length), float(stddev)))
+    residue_bonds["UNK"] = []
+
+    # Load bond angles.
+    residue_bond_angles: dict[str, list[BondAngle]] = {}
+    next(lines_iter)  # Skip empty line.
+    next(lines_iter)  # Skip header line.
+    for line in lines_iter:
+        if line.strip() == "-":
+            break
+        bond, resname, angle_degree, stddev_degree = line.split()
+        atom1, atom2, atom3 = bond.split("-")
+        if resname not in residue_bond_angles:
+            residue_bond_angles[resname] = []
+        residue_bond_angles[resname].append(
+            BondAngle(
+                atom1,
+                atom2,
+                atom3,
+                float(angle_degree) / 180.0 * np.pi,
+                float(stddev_degree) / 180.0 * np.pi,
+            )
+        )
+    residue_bond_angles["UNK"] = []
+
+    def make_bond_key(atom1_name: str, atom2_name: str) -> str:
+        """Unique key to lookup bonds."""
+        return "-".join(sorted([atom1_name, atom2_name]))
+
+    # Translate bond angles into distances ("virtual bonds").
+    residue_virtual_bonds: dict[str, list[Bond]] = {}
+    for resname, bond_angles in residue_bond_angles.items():
+        # Create a fast lookup dict for bond lengths.
+        bond_cache: dict[str, Bond] = {}
+        for b in residue_bonds[resname]:
+            bond_cache[make_bond_key(b.atom1_name, b.atom2_name)] = b
+        residue_virtual_bonds[resname] = []
+        for ba in bond_angles:
+            bond1 = bond_cache[make_bond_key(ba.atom1_name, ba.atom2_name)]
+            bond2 = bond_cache[make_bond_key(ba.atom2_name, ba.atom3name)]
+
+            # Compute distance between atom1 and atom3 using the law of cosines
+            # c^2 = a^2 + b^2 - 2ab*cos(gamma).
+            gamma = ba.angle_rad
+            length = np.sqrt(bond1.length**2 + bond2.length**2 - 2 * bond1.length * bond2.length * np.cos(gamma))
+
+            # Propagation of uncertainty assuming uncorrelated errors.
+            dl_outer = 0.5 / length
+            dl_dgamma = (2 * bond1.length * bond2.length * np.sin(gamma)) * dl_outer
+            dl_db1 = (2 * bond1.length - 2 * bond2.length * np.cos(gamma)) * dl_outer
+            dl_db2 = (2 * bond2.length - 2 * bond1.length * np.cos(gamma)) * dl_outer
+            stddev = np.sqrt(
+                (dl_dgamma * ba.stddev) ** 2 + (dl_db1 * bond1.stddev) ** 2 + (dl_db2 * bond2.stddev) ** 2
+            )
+            residue_virtual_bonds[resname].append(Bond(ba.atom1_name, ba.atom3name, length, stddev))
+
+    return (residue_bonds, residue_virtual_bonds, residue_bond_angles)
+
+
+# Between-residue bond lengths for general bonds (first element) and for Proline
+# (second element).
+between_res_bond_length_c_n: tuple[float, float] = (1.329, 1.341)
+between_res_bond_length_stddev_c_n: tuple[float, float] = (0.014, 0.016)
+
+# Between-residue cos_angles.
+between_res_cos_angles_c_n_ca: tuple[float, float] = (-0.5203, 0.0353)  # degrees: 121.352 +- 2.315
+between_res_cos_angles_ca_c_n: tuple[float, float] = (-0.4473, 0.0311)  # degrees: 116.568 +- 1.995
+
+# This mapping is used when we need to store atom data in a format that requires
+# fixed atom data size for every residue (e.g. a numpy array).
+atom_types: list[str] = [
+    "N",
+    "CA",
+    "C",
+    "CB",
+    "O",
+    "CG",
+    "CG1",
+    "CG2",
+    "OG",
+    "OG1",
+    "SG",
+    "CD",
+    "CD1",
+    "CD2",
+    "ND1",
+    "ND2",
+    "OD1",
+    "OD2",
+    "SD",
+    "CE",
+    "CE1",
+    "CE2",
+    "CE3",
+    "NE",
+    "NE1",
+    "NE2",
+    "OE1",
+    "OE2",
+    "CH2",
+    "NH1",
+    "NH2",
+    "OH",
+    "CZ",
+    "CZ2",
+    "CZ3",
+    "NZ",
+    "OXT",
+]
+atom_order: dict[str, int] = {atom_type: i for i, atom_type in enumerate(atom_types)}
+atom_type_num = len(atom_types)  # := 37.
+
+# A compact atom encoding with 14 columns
+# pylint: disable=line-too-long
+restype_name_to_atom14_names: dict[str, list[str]] = {
+    "ALA": ["N", "CA", "C", "O", "CB", "", "", "", "", "", "", "", "", ""],
+    "ARG": ["N", "CA", "C", "O", "CB", "CG", "CD", "NE", "CZ", "NH1", "NH2", "", "", ""],
+    "ASN": ["N", "CA", "C", "O", "CB", "CG", "OD1", "ND2", "", "", "", "", "", ""],
+    "ASP": ["N", "CA", "C", "O", "CB", "CG", "OD1", "OD2", "", "", "", "", "", ""],
+    "CYS": ["N", "CA", "C", "O", "CB", "SG", "", "", "", "", "", "", "", ""],
+    "GLN": ["N", "CA", "C", "O", "CB", "CG", "CD", "OE1", "NE2", "", "", "", "", ""],
+    "GLU": ["N", "CA", "C", "O", "CB", "CG", "CD", "OE1", "OE2", "", "", "", "", ""],
+    "GLY": ["N", "CA", "C", "O", "", "", "", "", "", "", "", "", "", ""],
+    "HIS": ["N", "CA", "C", "O", "CB", "CG", "ND1", "CD2", "CE1", "NE2", "", "", "", ""],
+    "ILE": ["N", "CA", "C", "O", "CB", "CG1", "CG2", "CD1", "", "", "", "", "", ""],
+    "LEU": ["N", "CA", "C", "O", "CB", "CG", "CD1", "CD2", "", "", "", "", "", ""],
+    "LYS": ["N", "CA", "C", "O", "CB", "CG", "CD", "CE", "NZ", "", "", "", "", ""],
+    "MET": ["N", "CA", "C", "O", "CB", "CG", "SD", "CE", "", "", "", "", "", ""],
+    "PHE": ["N", "CA", "C", "O", "CB", "CG", "CD1", "CD2", "CE1", "CE2", "CZ", "", "", ""],
+    "PRO": ["N", "CA", "C", "O", "CB", "CG", "CD", "", "", "", "", "", "", ""],
+    "SER": ["N", "CA", "C", "O", "CB", "OG", "", "", "", "", "", "", "", ""],
+    "THR": ["N", "CA", "C", "O", "CB", "OG1", "CG2", "", "", "", "", "", "", ""],
+    "TRP": ["N", "CA", "C", "O", "CB", "CG", "CD1", "CD2", "NE1", "CE2", "CE3", "CZ2", "CZ3", "CH2"],
+    "TYR": ["N", "CA", "C", "O", "CB", "CG", "CD1", "CD2", "CE1", "CE2", "CZ", "OH", "", ""],
+    "VAL": ["N", "CA", "C", "O", "CB", "CG1", "CG2", "", "", "", "", "", "", ""],
+    "UNK": ["", "", "", "", "", "", "", "", "", "", "", "", "", ""],
+}
+# pylint: enable=line-too-long
+
+
+# This is the standard residue order when coding AA type as a number.
+# Reproduce it by taking 3-letter AA codes and sorting them alphabetically.
+restypes: list[str] = [
+    "A",
+    "R",
+    "N",
+    "D",
+    "C",
+    "Q",
+    "E",
+    "G",
+    "H",
+    "I",
+    "L",
+    "K",
+    "M",
+    "F",
+    "P",
+    "S",
+    "T",
+    "W",
+    "Y",
+    "V",
+]
+restype_order: dict[str, int] = {restype: i for i, restype in enumerate(restypes)}
+restype_num = len(restypes)  # := 20.
+unk_restype_index = restype_num  # Catch-all index for unknown restypes.
+
+restypes_with_x: list[str] = restypes + ["X"]
+restype_order_with_x: dict[str, int] = {restype: i for i, restype in enumerate(restypes_with_x)}
+
+
+def sequence_to_onehot(sequence: str, mapping: Mapping[str, int], map_unknown_to_x: bool = False) -> np.ndarray:
+    """Maps the given sequence into a one-hot encoded matrix.
+
+    Args:
+      sequence: An amino acid sequence.
+      mapping: A dictionary mapping amino acids to integers.
+      map_unknown_to_x: If True, any amino acid that is not in the mapping will be
+        mapped to the unknown amino acid 'X'. If the mapping doesn't contain amino acid 'X', an error will be thrown.
+        If False, any amino acid not in the mapping will throw an error.
+
+    Returns:
+      A numpy array of shape (seq_len, num_unique_aas) with one-hot encoding of the sequence.
+
+    Raises:
+      ValueError: If the mapping doesn't contain values from 0 to
+        num_unique_aas - 1 without any gaps.
+    """
+    num_entries = max(mapping.values()) + 1
+
+    if sorted(set(mapping.values())) != list(range(num_entries)):
+        raise ValueError(
+            "The mapping must have values from 0 to num_unique_aas-1 without any gaps. Got: %s"
+            % sorted(mapping.values())
+        )
+
+    one_hot_arr = np.zeros((len(sequence), num_entries), dtype=np.int32)
+
+    for aa_index, aa_type in enumerate(sequence):
+        if map_unknown_to_x:
+            if aa_type.isalpha() and aa_type.isupper():
+                aa_id = mapping.get(aa_type, mapping["X"])
+            else:
+                raise ValueError(f"Invalid character in the sequence: {aa_type}")
+        else:
+            aa_id = mapping[aa_type]
+        one_hot_arr[aa_index, aa_id] = 1
+
+    return one_hot_arr
+
+
+restype_1to3: dict[str, str] = {
+    "A": "ALA",
+    "R": "ARG",
+    "N": "ASN",
+    "D": "ASP",
+    "C": "CYS",
+    "Q": "GLN",
+    "E": "GLU",
+    "G": "GLY",
+    "H": "HIS",
+    "I": "ILE",
+    "L": "LEU",
+    "K": "LYS",
+    "M": "MET",
+    "F": "PHE",
+    "P": "PRO",
+    "S": "SER",
+    "T": "THR",
+    "W": "TRP",
+    "Y": "TYR",
+    "V": "VAL",
+}
+
+
+# NB: restype_3to1 differs from Bio.PDB.protein_letters_3to1 by being a simple
+# 1-to-1 mapping of 3 letter names to one letter names. The latter contains
+# many more, and less common, three letter names as keys and maps many of these
+# to the same one letter name (including 'X' and 'U' which we don't use here).
+restype_3to1: dict[str, str] = {v: k for k, v in restype_1to3.items()}
+
+# Define a restype name for all unknown residues.
+unk_restype = "UNK"
+
+resnames: list[str] = [restype_1to3[r] for r in restypes] + [unk_restype]
+resname_to_idx: dict[str, int] = {resname: i for i, resname in enumerate(resnames)}
+
+
+# The mapping here uses hhblits convention, so that B is mapped to D, J and O
+# are mapped to X, U is mapped to C, and Z is mapped to E. Other than that the
+# remaining 20 amino acids are kept in alphabetical order.
+# There are 2 non-amino acid codes, X (representing any amino acid) and
+# "-" representing a missing amino acid in an alignment.  The id for these
+# codes is put at the end (20 and 21) so that they can easily be ignored if
+# desired.
+HHBLITS_AA_TO_ID: dict[str, int] = {
+    "A": 0,
+    "B": 2,
+    "C": 1,
+    "D": 2,
+    "E": 3,
+    "F": 4,
+    "G": 5,
+    "H": 6,
+    "I": 7,
+    "J": 20,
+    "K": 8,
+    "L": 9,
+    "M": 10,
+    "N": 11,
+    "O": 20,
+    "P": 12,
+    "Q": 13,
+    "R": 14,
+    "S": 15,
+    "T": 16,
+    "U": 1,
+    "V": 17,
+    "W": 18,
+    "X": 20,
+    "Y": 19,
+    "Z": 3,
+    "-": 21,
+}
+
+# Partial inversion of HHBLITS_AA_TO_ID.
+ID_TO_HHBLITS_AA: dict[int, str] = {
+    0: "A",
+    1: "C",  # Also U.
+    2: "D",  # Also B.
+    3: "E",  # Also Z.
+    4: "F",
+    5: "G",
+    6: "H",
+    7: "I",
+    8: "K",
+    9: "L",
+    10: "M",
+    11: "N",
+    12: "P",
+    13: "Q",
+    14: "R",
+    15: "S",
+    16: "T",
+    17: "V",
+    18: "W",
+    19: "Y",
+    20: "X",  # Includes J and O.
+    21: "-",
+}
+
+restypes_with_x_and_gap: list[str] = restypes + ["X", "-"]
+MAP_HHBLITS_AATYPE_TO_OUR_AATYPE: tuple[int, ...] = tuple(
+    restypes_with_x_and_gap.index(ID_TO_HHBLITS_AA[i]) for i in range(len(restypes_with_x_and_gap))
+)
+
+
+def _make_standard_atom_mask() -> np.ndarray:
+    """Returns [num_res_types, num_atom_types] mask array."""
+    # +1 to account for unknown (all 0s).
+    mask = np.zeros([restype_num + 1, atom_type_num], dtype=np.int32)
+    for restype, restype_letter in enumerate(restypes):
+        restype_name = restype_1to3[restype_letter]
+        atom_names = residue_atoms[restype_name]
+        for atom_name in atom_names:
+            atom_type = atom_order[atom_name]
+            mask[restype, atom_type] = 1
+    return mask
+
+
+STANDARD_ATOM_MASK = _make_standard_atom_mask()
+
+
+# A one hot representation for the first and second atoms defining the axis
+# of rotation for each chi-angle in each residue.
+def chi_angle_atom(atom_index: int) -> np.ndarray:
+    """Define chi-angle rigid groups via one-hot representations."""
+    chi_angles_index = {}
+    one_hots = []
+
+    for k, v in chi_angles_atoms.items():
+        indices = [atom_types.index(s[atom_index]) for s in v]
+        indices.extend([-1] * (4 - len(indices)))
+        chi_angles_index[k] = indices
+
+    for r in restypes:
+        res3 = restype_1to3[r]
+        one_hot = np.eye(atom_type_num)[chi_angles_index[res3]]
+        one_hots.append(one_hot)
+
+    one_hots.append(np.zeros([4, atom_type_num]))  # Add zeros for residue `X`.
+    one_hot = np.stack(one_hots, axis=0)
+    one_hot = np.transpose(one_hot, [0, 2, 1])
+
+    return one_hot
+
+
+chi_atom_1_one_hot = chi_angle_atom(1)
+chi_atom_2_one_hot = chi_angle_atom(2)
+
+# An array like chi_angles_atoms but using indices rather than names.
+chi_angles_atom_indices_list: list[list[list[str]]] = [chi_angles_atoms[restype_1to3[r]] for r in restypes]
+chi_angles_atom_indices_ours: list = map_structure_with_atom_order(chi_angles_atom_indices_list)
+chi_angles_atom_indices = np.array(
+    [chi_atoms + ([[0, 0, 0, 0]] * (4 - len(chi_atoms))) for chi_atoms in chi_angles_atom_indices_list]
+)
+
+# Mapping from (res_name, atom_name) pairs to the atom's chi group index
+# and atom index within that group.
+chi_groups_for_atom: dict[tuple[str, str], list[tuple[int, int]]] = collections.defaultdict(list)
+for res_name, chi_angle_atoms_for_res in chi_angles_atoms.items():
+    for chi_group_i, chi_group in enumerate(chi_angle_atoms_for_res):
+        for atom_i, atom in enumerate(chi_group):
+            chi_groups_for_atom[(res_name, atom)].append((chi_group_i, atom_i))
+chi_groups_for_atom = dict(chi_groups_for_atom)
+
+
+def _make_rigid_transformation_4x4(ex: np.ndarray, ey: np.ndarray, translation: np.ndarray) -> np.ndarray:
+    """Create a rigid 4x4 transformation matrix from two axes and transl."""
+    # Normalize ex.
+    ex_normalized = ex / np.linalg.norm(ex)
+
+    # make ey perpendicular to ex
+    ey_normalized = ey - np.dot(ey, ex_normalized) * ex_normalized
+    ey_normalized /= np.linalg.norm(ey_normalized)
+
+    # compute ez as cross product
+    eznorm = np.cross(ex_normalized, ey_normalized)
+    m = np.stack([ex_normalized, ey_normalized, eznorm, translation]).transpose()
+    m = np.concatenate([m, [[0.0, 0.0, 0.0, 1.0]]], axis=0)
+    return m
+
+
+# create an array with (restype, atomtype) --> rigid_group_idx
+# and an array with (restype, atomtype, coord) for the atom positions
+# and compute affine transformation matrices (4,4) from one rigid group to the
+# previous group
+restype_atom37_to_rigid_group = np.zeros([21, 37], dtype=int)
+restype_atom37_mask = np.zeros([21, 37], dtype=np.float32)
+restype_atom37_rigid_group_positions = np.zeros([21, 37, 3], dtype=np.float32)
+restype_atom14_to_rigid_group = np.zeros([21, 14], dtype=int)
+restype_atom14_mask = np.zeros([21, 14], dtype=np.float32)
+restype_atom14_rigid_group_positions = np.zeros([21, 14, 3], dtype=np.float32)
+restype_rigid_group_default_frame = np.zeros([21, 8, 4, 4], dtype=np.float32)
+
+
+def _make_rigid_group_constants() -> None:
+    """Fill the arrays above."""
+    for restype, restype_letter in enumerate(restypes):
+        resname = restype_1to3[restype_letter]
+        for atomname, group_idx, atom_position in rigid_group_atom_positions[resname]:
+            atomtype = atom_order[atomname]
+            restype_atom37_to_rigid_group[restype, atomtype] = group_idx
+            restype_atom37_mask[restype, atomtype] = 1
+            restype_atom37_rigid_group_positions[restype, atomtype, :] = atom_position
+
+            atom14idx = restype_name_to_atom14_names[resname].index(atomname)
+            restype_atom14_to_rigid_group[restype, atom14idx] = group_idx
+            restype_atom14_mask[restype, atom14idx] = 1
+            restype_atom14_rigid_group_positions[restype, atom14idx, :] = atom_position
+
+    for restype, restype_letter in enumerate(restypes):
+        resname = restype_1to3[restype_letter]
+        atom_positions: dict[str, np.ndarray] = {
+            name: np.array(pos) for name, _, pos in rigid_group_atom_positions[resname]
+        }
+
+        # backbone to backbone is the identity transform
+        restype_rigid_group_default_frame[restype, 0, :, :] = np.eye(4)
+
+        # pre-omega-frame to backbone (currently dummy identity matrix)
+        restype_rigid_group_default_frame[restype, 1, :, :] = np.eye(4)
+
+        # phi-frame to backbone
+        mat = _make_rigid_transformation_4x4(
+            ex=atom_positions["N"] - atom_positions["CA"],
+            ey=np.array([1.0, 0.0, 0.0]),
+            translation=atom_positions["N"],
+        )
+        restype_rigid_group_default_frame[restype, 2, :, :] = mat
+
+        # psi-frame to backbone
+        mat = _make_rigid_transformation_4x4(
+            ex=atom_positions["C"] - atom_positions["CA"],
+            ey=atom_positions["CA"] - atom_positions["N"],
+            translation=atom_positions["C"],
+        )
+        restype_rigid_group_default_frame[restype, 3, :, :] = mat
+
+        # chi1-frame to backbone
+        if chi_angles_mask[restype][0]:
+            base_atom_names = chi_angles_atoms[resname][0]
+            base_atom_positions = [atom_positions[name] for name in base_atom_names]
+            mat = _make_rigid_transformation_4x4(
+                ex=base_atom_positions[2] - base_atom_positions[1],
+                ey=base_atom_positions[0] - base_atom_positions[1],
+                translation=base_atom_positions[2],
+            )
+            restype_rigid_group_default_frame[restype, 4, :, :] = mat
+
+        # chi2-frame to chi1-frame
+        # chi3-frame to chi2-frame
+        # chi4-frame to chi3-frame
+        # luckily all rotation axes for the next frame start at (0,0,0) of the
+        # previous frame
+        for chi_idx in range(1, 4):
+            if chi_angles_mask[restype][chi_idx]:
+                axis_end_atom_name = chi_angles_atoms[resname][chi_idx][2]
+                axis_end_atom_position = atom_positions[axis_end_atom_name]
+                mat = _make_rigid_transformation_4x4(
+                    ex=axis_end_atom_position,
+                    ey=np.array([-1.0, 0.0, 0.0]),
+                    translation=axis_end_atom_position,
+                )
+                restype_rigid_group_default_frame[restype, 4 + chi_idx, :, :] = mat
+
+
+_make_rigid_group_constants()
+
+
+def make_atom14_dists_bounds(
+    overlap_tolerance: float = 1.5,
+    bond_length_tolerance_factor: int = 15,
+) -> dict[str, np.ndarray]:
+    """compute upper and lower bounds for bonds to assess violations."""
+    restype_atom14_bond_lower_bound = np.zeros([21, 14, 14], np.float32)
+    restype_atom14_bond_upper_bound = np.zeros([21, 14, 14], np.float32)
+    restype_atom14_bond_stddev = np.zeros([21, 14, 14], np.float32)
+    residue_bonds, residue_virtual_bonds, _ = load_stereo_chemical_props()
+    for restype, restype_letter in enumerate(restypes):
+        resname = restype_1to3[restype_letter]
+        atom_list = restype_name_to_atom14_names[resname]
+
+        # create lower and upper bounds for clashes
+        for atom1_idx, atom1_name in enumerate(atom_list):
+            if not atom1_name:
+                continue
+            atom1_radius = van_der_waals_radius[atom1_name[0]]
+            for atom2_idx, atom2_name in enumerate(atom_list):
+                if (not atom2_name) or atom1_idx == atom2_idx:
+                    continue
+                atom2_radius = van_der_waals_radius[atom2_name[0]]
+                lower = atom1_radius + atom2_radius - overlap_tolerance
+                upper = 1e10
+                restype_atom14_bond_lower_bound[restype, atom1_idx, atom2_idx] = lower
+                restype_atom14_bond_lower_bound[restype, atom2_idx, atom1_idx] = lower
+                restype_atom14_bond_upper_bound[restype, atom1_idx, atom2_idx] = upper
+                restype_atom14_bond_upper_bound[restype, atom2_idx, atom1_idx] = upper
+
+        # overwrite lower and upper bounds for bonds and angles
+        for b in residue_bonds[resname] + residue_virtual_bonds[resname]:
+            atom1_idx = atom_list.index(b.atom1_name)
+            atom2_idx = atom_list.index(b.atom2_name)
+            lower = b.length - bond_length_tolerance_factor * b.stddev
+            upper = b.length + bond_length_tolerance_factor * b.stddev
+            restype_atom14_bond_lower_bound[restype, atom1_idx, atom2_idx] = lower
+            restype_atom14_bond_lower_bound[restype, atom2_idx, atom1_idx] = lower
+            restype_atom14_bond_upper_bound[restype, atom1_idx, atom2_idx] = upper
+            restype_atom14_bond_upper_bound[restype, atom2_idx, atom1_idx] = upper
+            restype_atom14_bond_stddev[restype, atom1_idx, atom2_idx] = b.stddev
+            restype_atom14_bond_stddev[restype, atom2_idx, atom1_idx] = b.stddev
+    return {
+        "lower_bound": restype_atom14_bond_lower_bound,  # shape (21,14,14)
+        "upper_bound": restype_atom14_bond_upper_bound,  # shape (21,14,14)
+        "stddev": restype_atom14_bond_stddev,  # shape (21,14,14)
+    }
+
+
+restype_atom14_ambiguous_atoms = np.zeros((21, 14), dtype=np.float32)
+restype_atom14_ambiguous_atoms_swap_idx: np.ndarray = np.tile(np.arange(14, dtype=int), (21, 1))
+
+
+def _make_atom14_ambiguity_feats() -> None:
+    for res, pairs in residue_atom_renaming_swaps.items():
+        res_idx = restype_order[restype_3to1[res]]
+        for atom1, atom2 in pairs.items():
+            atom1_idx = restype_name_to_atom14_names[res].index(atom1)
+            atom2_idx = restype_name_to_atom14_names[res].index(atom2)
+            restype_atom14_ambiguous_atoms[res_idx, atom1_idx] = 1
+            restype_atom14_ambiguous_atoms[res_idx, atom2_idx] = 1
+            restype_atom14_ambiguous_atoms_swap_idx[res_idx, atom1_idx] = atom2_idx
+            restype_atom14_ambiguous_atoms_swap_idx[res_idx, atom2_idx] = atom1_idx
+
+
+_make_atom14_ambiguity_feats()
+
+
+def aatype_to_str_sequence(aatype: Sequence[int]) -> str:
+    return "".join([restypes_with_x[aatype[i]] for i in range(len(aatype))])

LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/esm/openfold_utils/rigid_utils.py

@@ -0,0 +1,1243 @@
+# Copyright 2021 AlQuraishi Laboratory
+# Copyright 2021 DeepMind Technologies Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from collections.abc import Callable, Sequence
+from functools import cache
+from typing import Any
+
+import numpy as np
+import torch
+
+
+def rot_matmul(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
+    """
+    Performs matrix multiplication of two rotation matrix tensors. Written out by hand to avoid AMP downcasting.
+
+    Args:
+        a: [*, 3, 3] left multiplicand
+        b: [*, 3, 3] right multiplicand
+    Returns:
+        The product ab
+    """
+
+    def row_mul(i: int) -> torch.Tensor:
+        return torch.stack(
+            [
+                a[..., i, 0] * b[..., 0, 0] + a[..., i, 1] * b[..., 1, 0] + a[..., i, 2] * b[..., 2, 0],
+                a[..., i, 0] * b[..., 0, 1] + a[..., i, 1] * b[..., 1, 1] + a[..., i, 2] * b[..., 2, 1],
+                a[..., i, 0] * b[..., 0, 2] + a[..., i, 1] * b[..., 1, 2] + a[..., i, 2] * b[..., 2, 2],
+            ],
+            dim=-1,
+        )
+
+    return torch.stack(
+        [
+            row_mul(0),
+            row_mul(1),
+            row_mul(2),
+        ],
+        dim=-2,
+    )
+
+
+def rot_vec_mul(r: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
+    """
+    Applies a rotation to a vector. Written out by hand to avoid transfer to avoid AMP downcasting.
+
+    Args:
+        r: [*, 3, 3] rotation matrices
+        t: [*, 3] coordinate tensors
+    Returns:
+        [*, 3] rotated coordinates
+    """
+    x, y, z = torch.unbind(t, dim=-1)
+    return torch.stack(
+        [
+            r[..., 0, 0] * x + r[..., 0, 1] * y + r[..., 0, 2] * z,
+            r[..., 1, 0] * x + r[..., 1, 1] * y + r[..., 1, 2] * z,
+            r[..., 2, 0] * x + r[..., 2, 1] * y + r[..., 2, 2] * z,
+        ],
+        dim=-1,
+    )
+
+
+@cache
+def identity_rot_mats(
+    batch_dims: tuple[int, ...],
+    dtype: torch.dtype | None = None,
+    device: torch.device | None = None,
+    requires_grad: bool = True,
+) -> torch.Tensor:
+    rots = torch.eye(3, dtype=dtype, device=device, requires_grad=requires_grad)
+    rots = rots.view(*((1,) * len(batch_dims)), 3, 3)
+    rots = rots.expand(*batch_dims, -1, -1)
+    rots = rots.contiguous()
+
+    return rots
+
+
+@cache
+def identity_trans(
+    batch_dims: tuple[int, ...],
+    dtype: torch.dtype | None = None,
+    device: torch.device | None = None,
+    requires_grad: bool = True,
+) -> torch.Tensor:
+    trans = torch.zeros((*batch_dims, 3), dtype=dtype, device=device, requires_grad=requires_grad)
+    return trans
+
+
+@cache
+def identity_quats(
+    batch_dims: tuple[int, ...],
+    dtype: torch.dtype | None = None,
+    device: torch.device | None = None,
+    requires_grad: bool = True,
+) -> torch.Tensor:
+    quat = torch.zeros((*batch_dims, 4), dtype=dtype, device=device, requires_grad=requires_grad)
+
+    with torch.no_grad():
+        quat[..., 0] = 1
+
+    return quat
+
+
+_quat_elements: list[str] = ["a", "b", "c", "d"]
+_qtr_keys: list[str] = [l1 + l2 for l1 in _quat_elements for l2 in _quat_elements]
+_qtr_ind_dict: dict[str, int] = {key: ind for ind, key in enumerate(_qtr_keys)}
+
+
+def _to_mat(pairs: list[tuple[str, int]]) -> np.ndarray:
+    mat = np.zeros((4, 4))
+    for key, value in pairs:
+        ind = _qtr_ind_dict[key]
+        mat[ind // 4][ind % 4] = value
+
+    return mat
+
+
+_QTR_MAT = np.zeros((4, 4, 3, 3))
+_QTR_MAT[..., 0, 0] = _to_mat([("aa", 1), ("bb", 1), ("cc", -1), ("dd", -1)])
+_QTR_MAT[..., 0, 1] = _to_mat([("bc", 2), ("ad", -2)])
+_QTR_MAT[..., 0, 2] = _to_mat([("bd", 2), ("ac", 2)])
+_QTR_MAT[..., 1, 0] = _to_mat([("bc", 2), ("ad", 2)])
+_QTR_MAT[..., 1, 1] = _to_mat([("aa", 1), ("bb", -1), ("cc", 1), ("dd", -1)])
+_QTR_MAT[..., 1, 2] = _to_mat([("cd", 2), ("ab", -2)])
+_QTR_MAT[..., 2, 0] = _to_mat([("bd", 2), ("ac", -2)])
+_QTR_MAT[..., 2, 1] = _to_mat([("cd", 2), ("ab", 2)])
+_QTR_MAT[..., 2, 2] = _to_mat([("aa", 1), ("bb", -1), ("cc", -1), ("dd", 1)])
+
+
+def quat_to_rot(quat: torch.Tensor) -> torch.Tensor:
+    """
+    Converts a quaternion to a rotation matrix.
+
+    Args:
+        quat: [*, 4] quaternions
+    Returns:
+        [*, 3, 3] rotation matrices
+    """
+    # [*, 4, 4]
+    quat = quat[..., None] * quat[..., None, :]
+
+    # [4, 4, 3, 3]
+    mat = _get_quat("_QTR_MAT", dtype=quat.dtype, device=quat.device)
+
+    # [*, 4, 4, 3, 3]
+    shaped_qtr_mat = mat.view((1,) * len(quat.shape[:-2]) + mat.shape)
+    quat = quat[..., None, None] * shaped_qtr_mat
+
+    # [*, 3, 3]
+    return torch.sum(quat, dim=(-3, -4))
+
+
+def rot_to_quat(rot: torch.Tensor) -> torch.Tensor:
+    if rot.shape[-2:] != (3, 3):
+        raise ValueError("Input rotation is incorrectly shaped")
+
+    [[xx, xy, xz], [yx, yy, yz], [zx, zy, zz]] = [[rot[..., i, j] for j in range(3)] for i in range(3)]
+
+    k = [
+        [
+            xx + yy + zz,
+            zy - yz,
+            xz - zx,
+            yx - xy,
+        ],
+        [
+            zy - yz,
+            xx - yy - zz,
+            xy + yx,
+            xz + zx,
+        ],
+        [
+            xz - zx,
+            xy + yx,
+            yy - xx - zz,
+            yz + zy,
+        ],
+        [
+            yx - xy,
+            xz + zx,
+            yz + zy,
+            zz - xx - yy,
+        ],
+    ]
+
+    _, vectors = torch.linalg.eigh((1.0 / 3.0) * torch.stack([torch.stack(t, dim=-1) for t in k], dim=-2))
+    return vectors[..., -1]
+
+
+_QUAT_MULTIPLY = np.zeros((4, 4, 4))
+_QUAT_MULTIPLY[:, :, 0] = [[1, 0, 0, 0], [0, -1, 0, 0], [0, 0, -1, 0], [0, 0, 0, -1]]
+
+_QUAT_MULTIPLY[:, :, 1] = [[0, 1, 0, 0], [1, 0, 0, 0], [0, 0, 0, 1], [0, 0, -1, 0]]
+
+_QUAT_MULTIPLY[:, :, 2] = [[0, 0, 1, 0], [0, 0, 0, -1], [1, 0, 0, 0], [0, 1, 0, 0]]
+
+_QUAT_MULTIPLY[:, :, 3] = [[0, 0, 0, 1], [0, 0, 1, 0], [0, -1, 0, 0], [1, 0, 0, 0]]
+
+_QUAT_MULTIPLY_BY_VEC = _QUAT_MULTIPLY[:, 1:, :]
+
+_CACHED_QUATS: dict[str, np.ndarray] = {
+    "_QTR_MAT": _QTR_MAT,
+    "_QUAT_MULTIPLY": _QUAT_MULTIPLY,
+    "_QUAT_MULTIPLY_BY_VEC": _QUAT_MULTIPLY_BY_VEC,
+}
+
+
+@cache
+def _get_quat(quat_key: str, dtype: torch.dtype, device: torch.device) -> torch.Tensor:
+    return torch.tensor(_CACHED_QUATS[quat_key], dtype=dtype, device=device)
+
+
+def quat_multiply(quat1: torch.Tensor, quat2: torch.Tensor) -> torch.Tensor:
+    """Multiply a quaternion by another quaternion."""
+    mat = _get_quat("_QUAT_MULTIPLY", dtype=quat1.dtype, device=quat1.device)
+    reshaped_mat = mat.view((1,) * len(quat1.shape[:-1]) + mat.shape)
+    return torch.sum(reshaped_mat * quat1[..., :, None, None] * quat2[..., None, :, None], dim=(-3, -2))
+
+
+def quat_multiply_by_vec(quat: torch.Tensor, vec: torch.Tensor) -> torch.Tensor:
+    """Multiply a quaternion by a pure-vector quaternion."""
+    mat = _get_quat("_QUAT_MULTIPLY_BY_VEC", dtype=quat.dtype, device=quat.device)
+    reshaped_mat = mat.view((1,) * len(quat.shape[:-1]) + mat.shape)
+    return torch.sum(reshaped_mat * quat[..., :, None, None] * vec[..., None, :, None], dim=(-3, -2))
+
+
+def invert_rot_mat(rot_mat: torch.Tensor) -> torch.Tensor:
+    return rot_mat.transpose(-1, -2)
+
+
+def invert_quat(quat: torch.Tensor) -> torch.Tensor:
+    quat_prime = quat.clone()
+    quat_prime[..., 1:] *= -1
+    inv = quat_prime / torch.sum(quat**2, dim=-1, keepdim=True)
+    return inv
+
+
+class Rotation:
+    """
+    A 3D rotation. Depending on how the object is initialized, the rotation is represented by either a rotation matrix
+    or a quaternion, though both formats are made available by helper functions. To simplify gradient computation, the
+    underlying format of the rotation cannot be changed in-place. Like Rigid, the class is designed to mimic the
+    behavior of a torch Tensor, almost as if each Rotation object were a tensor of rotations, in one format or another.
+    """
+
+    def __init__(
+        self,
+        rot_mats: torch.Tensor | None = None,
+        quats: torch.Tensor | None = None,
+        normalize_quats: bool = True,
+    ):
+        """
+        Args:
+            rot_mats:
+                A [*, 3, 3] rotation matrix tensor. Mutually exclusive with quats
+            quats:
+                A [*, 4] quaternion. Mutually exclusive with rot_mats. If normalize_quats is not True, must be a unit
+                quaternion
+            normalize_quats:
+                If quats is specified, whether to normalize quats
+        """
+        if (rot_mats is None and quats is None) or (rot_mats is not None and quats is not None):
+            raise ValueError("Exactly one input argument must be specified")
+
+        if (rot_mats is not None and rot_mats.shape[-2:] != (3, 3)) or (quats is not None and quats.shape[-1] != 4):
+            raise ValueError("Incorrectly shaped rotation matrix or quaternion")
+
+        # Force full-precision
+        if quats is not None:
+            quats = quats.to(dtype=torch.float32)
+        if rot_mats is not None:
+            rot_mats = rot_mats.to(dtype=torch.float32)
+
+        if quats is not None and normalize_quats:
+            quats = quats / torch.linalg.norm(quats, dim=-1, keepdim=True)
+
+        self._rot_mats = rot_mats
+        self._quats = quats
+
+    @staticmethod
+    def identity(
+        shape,
+        dtype: torch.dtype | None = None,
+        device: torch.device | None = None,
+        requires_grad: bool = True,
+        fmt: str = "quat",
+    ) -> Rotation:
+        """
+        Returns an identity Rotation.
+
+        Args:
+            shape:
+                The "shape" of the resulting Rotation object. See documentation for the shape property
+            dtype:
+                The torch dtype for the rotation
+            device:
+                The torch device for the new rotation
+            requires_grad:
+                Whether the underlying tensors in the new rotation object should require gradient computation
+            fmt:
+                One of "quat" or "rot_mat". Determines the underlying format of the new object's rotation
+        Returns:
+            A new identity rotation
+        """
+        if fmt == "rot_mat":
+            rot_mats = identity_rot_mats(
+                shape,
+                dtype,
+                device,
+                requires_grad,
+            )
+            return Rotation(rot_mats=rot_mats, quats=None)
+        elif fmt == "quat":
+            quats = identity_quats(shape, dtype, device, requires_grad)
+            return Rotation(rot_mats=None, quats=quats, normalize_quats=False)
+        else:
+            raise ValueError(f"Invalid format: f{fmt}")
+
+    # Magic methods
+
+    def __getitem__(self, index: Any) -> Rotation:
+        """
+        Allows torch-style indexing over the virtual shape of the rotation object. See documentation for the shape
+        property.
+
+        Args:
+            index:
+                A torch index. E.g. (1, 3, 2), or (slice(None,))
+        Returns:
+            The indexed rotation
+        """
+        if type(index) is not tuple:
+            index = (index,)
+
+        if self._rot_mats is not None:
+            rot_mats = self._rot_mats[index + (slice(None), slice(None))]
+            return Rotation(rot_mats=rot_mats)
+        elif self._quats is not None:
+            quats = self._quats[index + (slice(None),)]
+            return Rotation(quats=quats, normalize_quats=False)
+        else:
+            raise ValueError("Both rotations are None")
+
+    def __mul__(self, right: torch.Tensor) -> Rotation:
+        """
+        Pointwise left multiplication of the rotation with a tensor. Can be used to e.g. mask the Rotation.
+
+        Args:
+            right:
+                The tensor multiplicand
+        Returns:
+            The product
+        """
+        if not (isinstance(right, torch.Tensor)):
+            raise TypeError("The other multiplicand must be a Tensor")
+
+        if self._rot_mats is not None:
+            rot_mats = self._rot_mats * right[..., None, None]
+            return Rotation(rot_mats=rot_mats, quats=None)
+        elif self._quats is not None:
+            quats = self._quats * right[..., None]
+            return Rotation(rot_mats=None, quats=quats, normalize_quats=False)
+        else:
+            raise ValueError("Both rotations are None")
+
+    def __rmul__(self, left: torch.Tensor) -> Rotation:
+        """
+        Reverse pointwise multiplication of the rotation with a tensor.
+
+        Args:
+            left:
+                The left multiplicand
+        Returns:
+            The product
+        """
+        return self.__mul__(left)
+
+    # Properties
+
+    @property
+    def shape(self) -> torch.Size:
+        """
+        Returns the virtual shape of the rotation object. This shape is defined as the batch dimensions of the
+        underlying rotation matrix or quaternion. If the Rotation was initialized with a [10, 3, 3] rotation matrix
+        tensor, for example, the resulting shape would be [10].
+
+        Returns:
+            The virtual shape of the rotation object
+        """
+        if self._rot_mats is not None:
+            return self._rot_mats.shape[:-2]
+        elif self._quats is not None:
+            return self._quats.shape[:-1]
+        else:
+            raise ValueError("Both rotations are None")
+
+    @property
+    def dtype(self) -> torch.dtype:
+        """
+        Returns the dtype of the underlying rotation.
+
+        Returns:
+            The dtype of the underlying rotation
+        """
+        if self._rot_mats is not None:
+            return self._rot_mats.dtype
+        elif self._quats is not None:
+            return self._quats.dtype
+        else:
+            raise ValueError("Both rotations are None")
+
+    @property
+    def device(self) -> torch.device:
+        """
+        The device of the underlying rotation
+
+        Returns:
+            The device of the underlying rotation
+        """
+        if self._rot_mats is not None:
+            return self._rot_mats.device
+        elif self._quats is not None:
+            return self._quats.device
+        else:
+            raise ValueError("Both rotations are None")
+
+    @property
+    def requires_grad(self) -> bool:
+        """
+        Returns the requires_grad property of the underlying rotation
+
+        Returns:
+            The requires_grad property of the underlying tensor
+        """
+        if self._rot_mats is not None:
+            return self._rot_mats.requires_grad
+        elif self._quats is not None:
+            return self._quats.requires_grad
+        else:
+            raise ValueError("Both rotations are None")
+
+    def get_rot_mats(self) -> torch.Tensor:
+        """
+        Returns the underlying rotation as a rotation matrix tensor.
+
+        Returns:
+            The rotation as a rotation matrix tensor
+        """
+        if self._rot_mats is not None:
+            return self._rot_mats
+        elif self._quats is not None:
+            return quat_to_rot(self._quats)
+        else:
+            raise ValueError("Both rotations are None")
+
+    def get_quats(self) -> torch.Tensor:
+        """
+        Returns the underlying rotation as a quaternion tensor.
+
+        Depending on whether the Rotation was initialized with a quaternion, this function may call torch.linalg.eigh.
+
+        Returns:
+            The rotation as a quaternion tensor.
+        """
+        if self._rot_mats is not None:
+            return rot_to_quat(self._rot_mats)
+        elif self._quats is not None:
+            return self._quats
+        else:
+            raise ValueError("Both rotations are None")
+
+    def get_cur_rot(self) -> torch.Tensor:
+        """
+        Return the underlying rotation in its current form
+
+        Returns:
+            The stored rotation
+        """
+        if self._rot_mats is not None:
+            return self._rot_mats
+        elif self._quats is not None:
+            return self._quats
+        else:
+            raise ValueError("Both rotations are None")
+
+    # Rotation functions
+
+    def compose_q_update_vec(self, q_update_vec: torch.Tensor, normalize_quats: bool = True) -> Rotation:
+        """
+        Returns a new quaternion Rotation after updating the current object's underlying rotation with a quaternion
+        update, formatted as a [*, 3] tensor whose final three columns represent x, y, z such that (1, x, y, z) is the
+        desired (not necessarily unit) quaternion update.
+
+        Args:
+            q_update_vec:
+                A [*, 3] quaternion update tensor
+            normalize_quats:
+                Whether to normalize the output quaternion
+        Returns:
+            An updated Rotation
+        """
+        quats = self.get_quats()
+        new_quats = quats + quat_multiply_by_vec(quats, q_update_vec)
+        return Rotation(
+            rot_mats=None,
+            quats=new_quats,
+            normalize_quats=normalize_quats,
+        )
+
+    def compose_r(self, r: Rotation) -> Rotation:
+        """
+        Compose the rotation matrices of the current Rotation object with those of another.
+
+        Args:
+            r:
+                An update rotation object
+        Returns:
+            An updated rotation object
+        """
+        r1 = self.get_rot_mats()
+        r2 = r.get_rot_mats()
+        new_rot_mats = rot_matmul(r1, r2)
+        return Rotation(rot_mats=new_rot_mats, quats=None)
+
+    def compose_q(self, r: Rotation, normalize_quats: bool = True) -> Rotation:
+        """
+        Compose the quaternions of the current Rotation object with those of another.
+
+        Depending on whether either Rotation was initialized with quaternions, this function may call
+        torch.linalg.eigh.
+
+        Args:
+            r:
+                An update rotation object
+        Returns:
+            An updated rotation object
+        """
+        q1 = self.get_quats()
+        q2 = r.get_quats()
+        new_quats = quat_multiply(q1, q2)
+        return Rotation(rot_mats=None, quats=new_quats, normalize_quats=normalize_quats)
+
+    def apply(self, pts: torch.Tensor) -> torch.Tensor:
+        """
+        Apply the current Rotation as a rotation matrix to a set of 3D coordinates.
+
+        Args:
+            pts:
+                A [*, 3] set of points
+        Returns:
+            [*, 3] rotated points
+        """
+        rot_mats = self.get_rot_mats()
+        return rot_vec_mul(rot_mats, pts)
+
+    def invert_apply(self, pts: torch.Tensor) -> torch.Tensor:
+        """
+        The inverse of the apply() method.
+
+        Args:
+            pts:
+                A [*, 3] set of points
+        Returns:
+            [*, 3] inverse-rotated points
+        """
+        rot_mats = self.get_rot_mats()
+        inv_rot_mats = invert_rot_mat(rot_mats)
+        return rot_vec_mul(inv_rot_mats, pts)
+
+    def invert(self) -> Rotation:
+        """
+        Returns the inverse of the current Rotation.
+
+        Returns:
+            The inverse of the current Rotation
+        """
+        if self._rot_mats is not None:
+            return Rotation(rot_mats=invert_rot_mat(self._rot_mats), quats=None)
+        elif self._quats is not None:
+            return Rotation(
+                rot_mats=None,
+                quats=invert_quat(self._quats),
+                normalize_quats=False,
+            )
+        else:
+            raise ValueError("Both rotations are None")
+
+    # "Tensor" stuff
+
+    def unsqueeze(self, dim: int) -> Rotation:
+        """
+        Analogous to torch.unsqueeze. The dimension is relative to the shape of the Rotation object.
+
+        Args:
+            dim: A positive or negative dimension index.
+        Returns:
+            The unsqueezed Rotation.
+        """
+        if dim >= len(self.shape):
+            raise ValueError("Invalid dimension")
+
+        if self._rot_mats is not None:
+            rot_mats = self._rot_mats.unsqueeze(dim if dim >= 0 else dim - 2)
+            return Rotation(rot_mats=rot_mats, quats=None)
+        elif self._quats is not None:
+            quats = self._quats.unsqueeze(dim if dim >= 0 else dim - 1)
+            return Rotation(rot_mats=None, quats=quats, normalize_quats=False)
+        else:
+            raise ValueError("Both rotations are None")
+
+    @staticmethod
+    def cat(rs: Sequence[Rotation], dim: int) -> Rotation:
+        """
+        Concatenates rotations along one of the batch dimensions. Analogous to torch.cat().
+
+        Note that the output of this operation is always a rotation matrix, regardless of the format of input
+        rotations.
+
+        Args:
+            rs:
+                A list of rotation objects
+            dim:
+                The dimension along which the rotations should be concatenated
+        Returns:
+            A concatenated Rotation object in rotation matrix format
+        """
+        rot_mats = torch.cat(
+            [r.get_rot_mats() for r in rs],
+            dim=dim if dim >= 0 else dim - 2,
+        )
+
+        return Rotation(rot_mats=rot_mats, quats=None)
+
+    def map_tensor_fn(self, fn: Callable[[torch.Tensor], torch.Tensor]) -> Rotation:
+        """
+        Apply a Tensor -> Tensor function to underlying rotation tensors, mapping over the rotation dimension(s). Can
+        be used e.g. to sum out a one-hot batch dimension.
+
+        Args:
+            fn:
+                A Tensor -> Tensor function to be mapped over the Rotation
+        Returns:
+            The transformed Rotation object
+        """
+        if self._rot_mats is not None:
+            rot_mats = self._rot_mats.view(self._rot_mats.shape[:-2] + (9,))
+            rot_mats = torch.stack(list(map(fn, torch.unbind(rot_mats, dim=-1))), dim=-1)
+            rot_mats = rot_mats.view(rot_mats.shape[:-1] + (3, 3))
+            return Rotation(rot_mats=rot_mats, quats=None)
+        elif self._quats is not None:
+            quats = torch.stack(list(map(fn, torch.unbind(self._quats, dim=-1))), dim=-1)
+            return Rotation(rot_mats=None, quats=quats, normalize_quats=False)
+        else:
+            raise ValueError("Both rotations are None")
+
+    def cuda(self) -> Rotation:
+        """
+        Analogous to the cuda() method of torch Tensors
+
+        Returns:
+            A copy of the Rotation in CUDA memory
+        """
+        if self._rot_mats is not None:
+            return Rotation(rot_mats=self._rot_mats.cuda(), quats=None)
+        elif self._quats is not None:
+            return Rotation(rot_mats=None, quats=self._quats.cuda(), normalize_quats=False)
+        else:
+            raise ValueError("Both rotations are None")
+
+    def to(self, device: torch.device | None, dtype: torch.dtype | None) -> Rotation:
+        """
+        Analogous to the to() method of torch Tensors
+
+        Args:
+            device:
+                A torch device
+            dtype:
+                A torch dtype
+        Returns:
+            A copy of the Rotation using the new device and dtype
+        """
+        if self._rot_mats is not None:
+            return Rotation(
+                rot_mats=self._rot_mats.to(device=device, dtype=dtype),
+                quats=None,
+            )
+        elif self._quats is not None:
+            return Rotation(
+                rot_mats=None,
+                quats=self._quats.to(device=device, dtype=dtype),
+                normalize_quats=False,
+            )
+        else:
+            raise ValueError("Both rotations are None")
+
+    def detach(self) -> Rotation:
+        """
+        Returns a copy of the Rotation whose underlying Tensor has been detached from its torch graph.
+
+        Returns:
+            A copy of the Rotation whose underlying Tensor has been detached from its torch graph
+        """
+        if self._rot_mats is not None:
+            return Rotation(rot_mats=self._rot_mats.detach(), quats=None)
+        elif self._quats is not None:
+            return Rotation(
+                rot_mats=None,
+                quats=self._quats.detach(),
+                normalize_quats=False,
+            )
+        else:
+            raise ValueError("Both rotations are None")
+
+
+class Rigid:
+    """
+    A class representing a rigid transformation. Little more than a wrapper around two objects: a Rotation object and a
+    [*, 3] translation Designed to behave approximately like a single torch tensor with the shape of the shared batch
+    dimensions of its component parts.
+    """
+
+    def __init__(self, rots: Rotation | None, trans: torch.Tensor | None):
+        """
+        Args:
+            rots: A [*, 3, 3] rotation tensor
+            trans: A corresponding [*, 3] translation tensor
+        """
+        # (we need device, dtype, etc. from at least one input)
+
+        batch_dims, dtype, device, requires_grad = None, None, None, None
+        if trans is not None:
+            batch_dims = trans.shape[:-1]
+            dtype = trans.dtype
+            device = trans.device
+            requires_grad = trans.requires_grad
+        elif rots is not None:
+            batch_dims = rots.shape
+            dtype = rots.dtype
+            device = rots.device
+            requires_grad = rots.requires_grad
+        else:
+            raise ValueError("At least one input argument must be specified")
+
+        if rots is None:
+            rots = Rotation.identity(
+                batch_dims,
+                dtype,
+                device,
+                requires_grad,
+            )
+        elif trans is None:
+            trans = identity_trans(
+                batch_dims,
+                dtype,
+                device,
+                requires_grad,
+            )
+
+        assert rots is not None
+        assert trans is not None
+
+        if (rots.shape != trans.shape[:-1]) or (rots.device != trans.device):
+            raise ValueError("Rots and trans incompatible")
+
+        # Force full precision. Happens to the rotations automatically.
+        trans = trans.to(dtype=torch.float32)
+
+        self._rots = rots
+        self._trans = trans
+
+    @staticmethod
+    def identity(
+        shape: tuple[int, ...],
+        dtype: torch.dtype | None = None,
+        device: torch.device | None = None,
+        requires_grad: bool = True,
+        fmt: str = "quat",
+    ) -> Rigid:
+        """
+        Constructs an identity transformation.
+
+        Args:
+            shape:
+                The desired shape
+            dtype:
+                The dtype of both internal tensors
+            device:
+                The device of both internal tensors
+            requires_grad:
+                Whether grad should be enabled for the internal tensors
+        Returns:
+            The identity transformation
+        """
+        return Rigid(
+            Rotation.identity(shape, dtype, device, requires_grad, fmt=fmt),
+            identity_trans(shape, dtype, device, requires_grad),
+        )
+
+    def __getitem__(self, index: Any) -> Rigid:
+        """
+        Indexes the affine transformation with PyTorch-style indices. The index is applied to the shared dimensions of
+        both the rotation and the translation.
+
+        E.g.::
+
+            r = Rotation(rot_mats=torch.rand(10, 10, 3, 3), quats=None) t = Rigid(r, torch.rand(10, 10, 3)) indexed =
+            t[3, 4:6] assert(indexed.shape == (2,)) assert(indexed.get_rots().shape == (2,))
+            assert(indexed.get_trans().shape == (2, 3))
+
+        Args:
+            index: A standard torch tensor index. E.g. 8, (10, None, 3),
+            or (3, slice(0, 1, None))
+        Returns:
+            The indexed tensor
+        """
+        if type(index) is not tuple:
+            index = (index,)
+
+        return Rigid(
+            self._rots[index],
+            self._trans[index + (slice(None),)],
+        )
+
+    def __mul__(self, right: torch.Tensor) -> Rigid:
+        """
+        Pointwise left multiplication of the transformation with a tensor. Can be used to e.g. mask the Rigid.
+
+        Args:
+            right:
+                The tensor multiplicand
+        Returns:
+            The product
+        """
+        if not (isinstance(right, torch.Tensor)):
+            raise TypeError("The other multiplicand must be a Tensor")
+
+        new_rots = self._rots * right
+        new_trans = self._trans * right[..., None]
+
+        return Rigid(new_rots, new_trans)
+
+    def __rmul__(self, left: torch.Tensor) -> Rigid:
+        """
+        Reverse pointwise multiplication of the transformation with a tensor.
+
+        Args:
+            left:
+                The left multiplicand
+        Returns:
+            The product
+        """
+        return self.__mul__(left)
+
+    @property
+    def shape(self) -> torch.Size:
+        """
+        Returns the shape of the shared dimensions of the rotation and the translation.
+
+        Returns:
+            The shape of the transformation
+        """
+        return self._trans.shape[:-1]
+
+    @property
+    def device(self) -> torch.device:
+        """
+        Returns the device on which the Rigid's tensors are located.
+
+        Returns:
+            The device on which the Rigid's tensors are located
+        """
+        return self._trans.device
+
+    def get_rots(self) -> Rotation:
+        """
+        Getter for the rotation.
+
+        Returns:
+            The rotation object
+        """
+        return self._rots
+
+    def get_trans(self) -> torch.Tensor:
+        """
+        Getter for the translation.
+
+        Returns:
+            The stored translation
+        """
+        return self._trans
+
+    def compose_q_update_vec(self, q_update_vec: torch.Tensor) -> Rigid:
+        """
+        Composes the transformation with a quaternion update vector of shape [*, 6], where the final 6 columns
+        represent the x, y, and z values of a quaternion of form (1, x, y, z) followed by a 3D translation.
+
+        Args:
+            q_vec: The quaternion update vector.
+        Returns:
+            The composed transformation.
+        """
+        q_vec, t_vec = q_update_vec[..., :3], q_update_vec[..., 3:]
+        new_rots = self._rots.compose_q_update_vec(q_vec)
+
+        trans_update = self._rots.apply(t_vec)
+        new_translation = self._trans + trans_update
+
+        return Rigid(new_rots, new_translation)
+
+    def compose(self, r: Rigid) -> Rigid:
+        """
+        Composes the current rigid object with another.
+
+        Args:
+            r:
+                Another Rigid object
+        Returns:
+            The composition of the two transformations
+        """
+        new_rot = self._rots.compose_r(r._rots)
+        new_trans = self._rots.apply(r._trans) + self._trans
+        return Rigid(new_rot, new_trans)
+
+    def apply(self, pts: torch.Tensor) -> torch.Tensor:
+        """
+        Applies the transformation to a coordinate tensor.
+
+        Args:
+            pts: A [*, 3] coordinate tensor.
+        Returns:
+            The transformed points.
+        """
+        rotated = self._rots.apply(pts)
+        return rotated + self._trans
+
+    def invert_apply(self, pts: torch.Tensor) -> torch.Tensor:
+        """
+        Applies the inverse of the transformation to a coordinate tensor.
+
+        Args:
+            pts: A [*, 3] coordinate tensor
+        Returns:
+            The transformed points.
+        """
+        pts = pts - self._trans
+        return self._rots.invert_apply(pts)
+
+    def invert(self) -> Rigid:
+        """
+        Inverts the transformation.
+
+        Returns:
+            The inverse transformation.
+        """
+        rot_inv = self._rots.invert()
+        trn_inv = rot_inv.apply(self._trans)
+
+        return Rigid(rot_inv, -1 * trn_inv)
+
+    def map_tensor_fn(self, fn: Callable[[torch.Tensor], torch.Tensor]) -> Rigid:
+        """
+        Apply a Tensor -> Tensor function to underlying translation and rotation tensors, mapping over the
+        translation/rotation dimensions respectively.
+
+        Args:
+            fn:
+                A Tensor -> Tensor function to be mapped over the Rigid
+        Returns:
+            The transformed Rigid object
+        """
+        new_rots = self._rots.map_tensor_fn(fn)
+        new_trans = torch.stack(list(map(fn, torch.unbind(self._trans, dim=-1))), dim=-1)
+
+        return Rigid(new_rots, new_trans)
+
+    def to_tensor_4x4(self) -> torch.Tensor:
+        """
+        Converts a transformation to a homogeneous transformation tensor.
+
+        Returns:
+            A [*, 4, 4] homogeneous transformation tensor
+        """
+        tensor = self._trans.new_zeros((*self.shape, 4, 4))
+        tensor[..., :3, :3] = self._rots.get_rot_mats()
+        tensor[..., :3, 3] = self._trans
+        tensor[..., 3, 3] = 1
+        return tensor
+
+    @staticmethod
+    def from_tensor_4x4(t: torch.Tensor) -> Rigid:
+        """
+        Constructs a transformation from a homogeneous transformation tensor.
+
+        Args:
+            t: [*, 4, 4] homogeneous transformation tensor
+        Returns:
+            T object with shape [*]
+        """
+        if t.shape[-2:] != (4, 4):
+            raise ValueError("Incorrectly shaped input tensor")
+
+        rots = Rotation(rot_mats=t[..., :3, :3], quats=None)
+        trans = t[..., :3, 3]
+
+        return Rigid(rots, trans)
+
+    def to_tensor_7(self) -> torch.Tensor:
+        """
+        Converts a transformation to a tensor with 7 final columns, four for the quaternion followed by three for the
+        translation.
+
+        Returns:
+            A [*, 7] tensor representation of the transformation
+        """
+        tensor = self._trans.new_zeros((*self.shape, 7))
+        tensor[..., :4] = self._rots.get_quats()
+        tensor[..., 4:] = self._trans
+
+        return tensor
+
+    @staticmethod
+    def from_tensor_7(t: torch.Tensor, normalize_quats: bool = False) -> Rigid:
+        if t.shape[-1] != 7:
+            raise ValueError("Incorrectly shaped input tensor")
+
+        quats, trans = t[..., :4], t[..., 4:]
+
+        rots = Rotation(rot_mats=None, quats=quats, normalize_quats=normalize_quats)
+
+        return Rigid(rots, trans)
+
+    @staticmethod
+    def from_3_points(
+        p_neg_x_axis: torch.Tensor, origin: torch.Tensor, p_xy_plane: torch.Tensor, eps: float = 1e-8
+    ) -> Rigid:
+        """
+        Implements algorithm 21. Constructs transformations from sets of 3 points using the Gram-Schmidt algorithm.
+
+        Args:
+            p_neg_x_axis: [*, 3] coordinates
+            origin: [*, 3] coordinates used as frame origins
+            p_xy_plane: [*, 3] coordinates
+            eps: Small epsilon value
+        Returns:
+            A transformation object of shape [*]
+        """
+        p_neg_x_axis_unbound = torch.unbind(p_neg_x_axis, dim=-1)
+        origin_unbound = torch.unbind(origin, dim=-1)
+        p_xy_plane_unbound = torch.unbind(p_xy_plane, dim=-1)
+
+        e0 = [c1 - c2 for c1, c2 in zip(origin_unbound, p_neg_x_axis_unbound)]
+        e1 = [c1 - c2 for c1, c2 in zip(p_xy_plane_unbound, origin_unbound)]
+
+        denom = torch.sqrt(sum(c * c for c in e0) + eps * torch.ones_like(e0[0]))
+        e0 = [c / denom for c in e0]
+        dot = sum((c1 * c2 for c1, c2 in zip(e0, e1)))
+        e1 = [c2 - c1 * dot for c1, c2 in zip(e0, e1)]
+        denom = torch.sqrt(sum(c * c for c in e1) + eps * torch.ones_like(e1[0]))
+        e1 = [c / denom for c in e1]
+        e2 = [
+            e0[1] * e1[2] - e0[2] * e1[1],
+            e0[2] * e1[0] - e0[0] * e1[2],
+            e0[0] * e1[1] - e0[1] * e1[0],
+        ]
+
+        rots = torch.stack([c for tup in zip(e0, e1, e2) for c in tup], dim=-1)
+        rots = rots.reshape(rots.shape[:-1] + (3, 3))
+
+        rot_obj = Rotation(rot_mats=rots, quats=None)
+
+        return Rigid(rot_obj, torch.stack(origin_unbound, dim=-1))
+
+    def unsqueeze(self, dim: int) -> Rigid:
+        """
+        Analogous to torch.unsqueeze. The dimension is relative to the shared dimensions of the rotation/translation.
+
+        Args:
+            dim: A positive or negative dimension index.
+        Returns:
+            The unsqueezed transformation.
+        """
+        if dim >= len(self.shape):
+            raise ValueError("Invalid dimension")
+        rots = self._rots.unsqueeze(dim)
+        trans = self._trans.unsqueeze(dim if dim >= 0 else dim - 1)
+
+        return Rigid(rots, trans)
+
+    @staticmethod
+    def cat(ts: Sequence[Rigid], dim: int) -> Rigid:
+        """
+        Concatenates transformations along a new dimension.
+
+        Args:
+            ts:
+                A list of T objects
+            dim:
+                The dimension along which the transformations should be concatenated
+        Returns:
+            A concatenated transformation object
+        """
+        rots = Rotation.cat([t._rots for t in ts], dim)
+        trans = torch.cat([t._trans for t in ts], dim=dim if dim >= 0 else dim - 1)
+
+        return Rigid(rots, trans)
+
+    def apply_rot_fn(self, fn: Callable[[Rotation], Rotation]) -> Rigid:
+        """
+        Applies a Rotation -> Rotation function to the stored rotation object.
+
+        Args:
+            fn: A function of type Rotation -> Rotation
+        Returns:
+            A transformation object with a transformed rotation.
+        """
+        return Rigid(fn(self._rots), self._trans)
+
+    def apply_trans_fn(self, fn: Callable[[torch.Tensor], torch.Tensor]) -> Rigid:
+        """
+        Applies a Tensor -> Tensor function to the stored translation.
+
+        Args:
+            fn:
+                A function of type Tensor -> Tensor to be applied to the translation
+        Returns:
+            A transformation object with a transformed translation.
+        """
+        return Rigid(self._rots, fn(self._trans))
+
+    def scale_translation(self, trans_scale_factor: float) -> Rigid:
+        """
+        Scales the translation by a constant factor.
+
+        Args:
+            trans_scale_factor:
+                The constant factor
+        Returns:
+            A transformation object with a scaled translation.
+        """
+        return self.apply_trans_fn(lambda t: t * trans_scale_factor)
+
+    def stop_rot_gradient(self) -> Rigid:
+        """
+        Detaches the underlying rotation object
+
+        Returns:
+            A transformation object with detached rotations
+        """
+        return self.apply_rot_fn(lambda r: r.detach())
+
+    @staticmethod
+    def make_transform_from_reference(
+        n_xyz: torch.Tensor, ca_xyz: torch.Tensor, c_xyz: torch.Tensor, eps: float = 1e-20
+    ) -> Rigid:
+        """
+        Returns a transformation object from reference coordinates.
+
+        Note that this method does not take care of symmetries. If you provide the atom positions in the non-standard
+        way, the N atom will end up not at [-0.527250, 1.359329, 0.0] but instead at [-0.527250, -1.359329, 0.0]. You
+        need to take care of such cases in your code.
+
+        Args:
+            n_xyz: A [*, 3] tensor of nitrogen xyz coordinates.
+            ca_xyz: A [*, 3] tensor of carbon alpha xyz coordinates.
+            c_xyz: A [*, 3] tensor of carbon xyz coordinates.
+        Returns:
+            A transformation object. After applying the translation and rotation to the reference backbone, the
+            coordinates will approximately equal to the input coordinates.
+        """
+        translation = -1 * ca_xyz
+        n_xyz = n_xyz + translation
+        c_xyz = c_xyz + translation
+
+        c_x, c_y, c_z = [c_xyz[..., i] for i in range(3)]
+        norm = torch.sqrt(eps + c_x**2 + c_y**2)
+        sin_c1 = -c_y / norm
+        cos_c1 = c_x / norm
+
+        c1_rots = sin_c1.new_zeros((*sin_c1.shape, 3, 3))
+        c1_rots[..., 0, 0] = cos_c1
+        c1_rots[..., 0, 1] = -1 * sin_c1
+        c1_rots[..., 1, 0] = sin_c1
+        c1_rots[..., 1, 1] = cos_c1
+        c1_rots[..., 2, 2] = 1
+
+        norm = torch.sqrt(eps + c_x**2 + c_y**2 + c_z**2)
+        sin_c2 = c_z / norm
+        cos_c2 = torch.sqrt(c_x**2 + c_y**2) / norm
+
+        c2_rots = sin_c2.new_zeros((*sin_c2.shape, 3, 3))
+        c2_rots[..., 0, 0] = cos_c2
+        c2_rots[..., 0, 2] = sin_c2
+        c2_rots[..., 1, 1] = 1
+        c2_rots[..., 2, 0] = -1 * sin_c2
+        c2_rots[..., 2, 2] = cos_c2
+
+        c_rots = rot_matmul(c2_rots, c1_rots)
+        n_xyz = rot_vec_mul(c_rots, n_xyz)
+
+        _, n_y, n_z = [n_xyz[..., i] for i in range(3)]
+        norm = torch.sqrt(eps + n_y**2 + n_z**2)
+        sin_n = -n_z / norm
+        cos_n = n_y / norm
+
+        n_rots = sin_c2.new_zeros((*sin_c2.shape, 3, 3))
+        n_rots[..., 0, 0] = 1
+        n_rots[..., 1, 1] = cos_n
+        n_rots[..., 1, 2] = -1 * sin_n
+        n_rots[..., 2, 1] = sin_n
+        n_rots[..., 2, 2] = cos_n
+
+        rots = rot_matmul(n_rots, c_rots)
+
+        rots = rots.transpose(-1, -2)
+        translation = -1 * translation
+
+        rot_obj = Rotation(rot_mats=rots, quats=None)
+
+        return Rigid(rot_obj, translation)
+
+    def cuda(self) -> Rigid:
+        """
+        Moves the transformation object to GPU memory
+
+        Returns:
+            A version of the transformation on GPU
+        """
+        return Rigid(self._rots.cuda(), self._trans.cuda())

LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/ovis2/modeling_ovis2.py

@@ -0,0 +1,730 @@
+#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+#           This file was automatically generated from src/transformers/models/ovis2/modular_ovis2.py.
+#               Do NOT edit this file manually as any edits will be overwritten by the generation of
+#             the file from the modular. If any change should be done, please apply the change to the
+#                          modular_ovis2.py file directly. One of our CI enforces this.
+#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import math
+from collections.abc import Callable
+from dataclasses import dataclass
+
+import torch
+from torch import nn
+
+from ... import initialization as init
+from ...activations import ACT2FN
+from ...cache_utils import Cache
+from ...generation import GenerationMixin
+from ...integrations import use_kernel_forward_from_hub
+from ...modeling_layers import GradientCheckpointingLayer
+from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPast, BaseModelOutputWithPooling
+from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
+from ...processing_utils import Unpack
+from ...utils import ModelOutput, TransformersKwargs, auto_docstring, can_return_tuple, torch_compilable_check
+from ...utils.generic import merge_with_config_defaults
+from ...utils.output_capturing import capture_outputs
+from ..auto import AutoModel
+from .configuration_ovis2 import Ovis2Config, Ovis2VisionConfig
+
+
+@auto_docstring
+@dataclass
+class BaseModelOutputWithVisualIndicatorFeatures(BaseModelOutputWithPooling):
+    r"""
+    visual_indicator_features (`torch.FloatTensor` of shape `(batch_size, visual_indicator_size)`):
+        Visual indicator features extracted from the model, which can be used for auxiliary tasks or further processing.
+    """
+
+    visual_indicator_features: torch.FloatTensor | None = None
+
+
+@auto_docstring(
+    custom_intro="""
+    Base class for Llava outputs, with hidden states and attentions.
+    """
+)
+@dataclass
+class Ovis2ModelOutputWithPast(BaseModelOutputWithPast):
+    r"""
+    past_key_values (`Cache`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
+        It is a [`~cache_utils.Cache`] instance. For more details, see our [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache).
+
+        Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see
+        `past_key_values` input) to speed up sequential decoding.
+    image_hidden_states (`torch.FloatTensor`, *optional*):
+        A `torch.FloatTensor` of size `(batch_size, num_images, sequence_length, hidden_size)`.
+        image_hidden_states of the model produced by the vision encoder and after projecting the last hidden state.
+    """
+
+    image_hidden_states: torch.FloatTensor | None = None
+
+
+@auto_docstring(
+    custom_intro="""
+    Base class for Ovis2 causal language model (or autoregressive) outputs.
+    """
+)
+@dataclass
+class Ovis2CausalLMOutputWithPast(ModelOutput):
+    r"""
+    loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):
+        Language modeling loss (for next-token prediction).
+    logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`):
+        Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
+    past_key_values (`Cache`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
+        It is a [`~cache_utils.Cache`] instance. For more details, see our [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache).
+
+        Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see
+        `past_key_values` input) to speed up sequential decoding.
+    image_hidden_states (`torch.FloatTensor`, *optional*):
+        A `torch.FloatTensor` of size (batch_size * num_patches, num_images, sequence_length, hidden_size)`.
+        image_hidden_states of the model produced by the vision encoder and after projecting the last hidden state.
+    """
+
+    loss: torch.FloatTensor | None = None
+    logits: torch.FloatTensor | None = None
+    past_key_values: Cache | None = None
+    hidden_states: tuple[torch.FloatTensor] | None = None
+    attentions: tuple[torch.FloatTensor] | None = None
+    image_hidden_states: torch.FloatTensor | None = None
+
+
+@use_kernel_forward_from_hub("RMSNorm")
+class Ovis2RMSNorm(nn.Module):
+    def __init__(self, hidden_size, eps: float = 1e-6) -> None:
+        """
+        Ovis2RMSNorm is equivalent to T5LayerNorm
+        """
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(torch.float32)
+        variance = hidden_states.pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+        return self.weight * hidden_states.to(input_dtype)
+
+    def extra_repr(self):
+        return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
+
+
+class Ovis2VisionMLP(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.hidden_size = config.hidden_size
+        self.intermediate_size = config.intermediate_size
+        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.mlp_bias)
+        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.mlp_bias)
+        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=config.mlp_bias)
+        self.act_fn = ACT2FN[config.hidden_act]
+
+    def forward(self, x):
+        down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+        return down_proj
+
+
+class Ovis2VisionEmbeddings(nn.Module):
+    def __init__(self, config: Ovis2VisionConfig):
+        super().__init__()
+        self.config = config
+        self.embed_dim = config.hidden_size
+        self.image_size = config.image_size
+        self.patch_size = config.patch_size
+
+        self.patch_embedding = nn.Conv2d(
+            in_channels=config.num_channels,
+            out_channels=self.embed_dim,
+            kernel_size=self.patch_size,
+            stride=self.patch_size,
+            padding="valid",
+        )
+
+        self.num_patches = (self.image_size // self.patch_size) ** 2
+        self.num_positions = self.num_patches
+        self.position_embedding = nn.Embedding(self.num_positions, self.embed_dim)
+        self.register_buffer("position_ids", torch.arange(self.num_positions).expand((1, -1)), persistent=False)
+        self.rms_norm = Ovis2RMSNorm(config.hidden_size, config.rms_norm_eps)
+
+    def forward(self, pixel_values: torch.FloatTensor) -> torch.Tensor:
+        target_dtype = self.patch_embedding.weight.dtype
+        patch_embeds = self.patch_embedding(pixel_values.to(dtype=target_dtype))
+        embeddings = patch_embeds.flatten(2).transpose(1, 2)
+        embeddings = self.rms_norm(embeddings)
+
+        embeddings = embeddings + self.position_embedding(self.position_ids)
+
+        return embeddings
+
+
+def eager_attention_forward(
+    module: nn.Module,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    attention_mask: torch.Tensor | None,
+    scaling: float,
+    dropout: float = 0.0,
+    **kwargs,
+):
+    attn_weights = torch.matmul(query, key.transpose(-1, -2)) * scaling
+    if attention_mask is not None:
+        attn_weights = attn_weights + attention_mask
+
+    attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
+    attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
+
+    attn_output = torch.matmul(attn_weights, value)
+    attn_output = attn_output.transpose(1, 2).contiguous()
+
+    return attn_output, attn_weights
+
+
+class Ovis2VisionAttention(nn.Module):
+    """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.embed_dim = config.hidden_size
+        self.num_heads = config.num_attention_heads
+        self.head_dim = self.embed_dim // self.num_heads
+        if self.head_dim * self.num_heads != self.embed_dim:
+            raise ValueError(
+                f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim} and `num_heads`:"
+                f" {self.num_heads})."
+            )
+        self.scale = self.head_dim**-0.5
+        self.dropout = config.attention_dropout
+        self.is_causal = False
+        self.k_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=config.qkv_bias)
+        self.v_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=config.qkv_bias)
+        self.q_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=config.qkv_bias)
+        self.out_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=config.qkv_bias)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        **kwargs,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        """Input shape: Batch x Time x Channel"""
+
+        input_shape = hidden_states.shape[:-1]
+
+        hidden_shape = (*input_shape, -1, self.head_dim)
+        queries = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        keys = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        values = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+
+        attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
+            self.config._attn_implementation, eager_attention_forward
+        )
+
+        attn_output, attn_weights = attention_interface(
+            self,
+            queries,
+            keys,
+            values,
+            attention_mask,
+            is_causal=self.is_causal,
+            scaling=self.scale,
+            dropout=0.0 if not self.training else self.dropout,
+        )
+
+        attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+        attn_output = self.out_proj(attn_output)
+
+        return attn_output, attn_weights
+
+
+class Ovis2MLP(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.hidden_size = config.hidden_size
+        self.intermediate_size = config.intermediate_size
+        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.mlp_bias)
+        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.mlp_bias)
+        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=config.mlp_bias)
+        self.act_fn = ACT2FN[config.hidden_act]
+
+    def forward(self, x):
+        down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+        return down_proj
+
+
+class Ovis2VisionEncoderLayer(GradientCheckpointingLayer):
+    def __init__(self, config: Ovis2VisionConfig):
+        super().__init__()
+        self.attention = Ovis2VisionAttention(config)
+        self.ffn = Ovis2MLP(config)
+        self.rms_norm1 = Ovis2RMSNorm(config.hidden_size, config.rms_norm_eps)
+        self.rms_norm2 = Ovis2RMSNorm(config.hidden_size, config.rms_norm_eps)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> torch.Tensor:
+        norm_hidden_states = self.rms_norm1(hidden_states)
+        attn_output, _ = self.attention(hidden_states=norm_hidden_states, attention_mask=attention_mask, **kwargs)
+
+        hidden_states = hidden_states + attn_output
+        norm_hidden_states = self.rms_norm2(hidden_states)
+        mlp_output = self.ffn(norm_hidden_states)
+
+        hidden_states = hidden_states + mlp_output
+        return hidden_states
+
+
+class Ovis2VisionEncoder(nn.Module):
+    """
+    Transformer encoder consisting of `config.num_hidden_layers` self attention layers. Each layer is a
+    [`Ovis2VisionEncoderLayer`].
+
+    Args:
+        config: Ovis2VisionConfig
+    """
+
+    def __init__(self, config: Ovis2VisionConfig):
+        super().__init__()
+        self.config = config
+        self.layers = nn.ModuleList([Ovis2VisionEncoderLayer(config) for _ in range(config.num_hidden_layers)])
+        self.gradient_checkpointing = False
+
+    # Ignore copy
+    @can_return_tuple
+    @auto_docstring
+    def forward(
+        self,
+        inputs_embeds,
+        attention_mask: torch.Tensor | None = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> BaseModelOutput:
+        hidden_states = inputs_embeds
+        for encoder_layer in self.layers:
+            hidden_states = encoder_layer(hidden_states, attention_mask, **kwargs)
+
+        return BaseModelOutput(last_hidden_state=hidden_states)
+
+
+class Ovis2VisionTransformer(nn.Module):
+    def __init__(self, config: Ovis2VisionConfig):
+        super().__init__()
+        self.config = config
+        self.embeddings = Ovis2VisionEmbeddings(config)
+        self.encoder = Ovis2VisionEncoder(config)
+        self.rms_norm = Ovis2RMSNorm(config.hidden_size, config.rms_norm_eps)
+        self.gradient_checkpointing = False
+
+    @can_return_tuple
+    def forward(
+        self,
+        pixel_values,
+        attention_mask: torch.Tensor | None = None,
+        **kwargs,
+    ):
+        hidden_states = self.embeddings(pixel_values)
+
+        encoder_outputs: BaseModelOutput = self.encoder(
+            inputs_embeds=hidden_states,
+            attention_mask=attention_mask,
+            **kwargs,
+        )
+
+        last_hidden_state = encoder_outputs.last_hidden_state
+        last_hidden_state = self.rms_norm(last_hidden_state)
+
+        return BaseModelOutput(last_hidden_state=last_hidden_state)
+
+
+class Ovis2VisualEmbeddingTable(nn.Embedding):
+    def forward(self, visual_tokens: torch.Tensor) -> torch.Tensor:
+        if visual_tokens.dtype in [torch.int8, torch.int16, torch.int32, torch.int64, torch.long]:
+            return super().forward(visual_tokens)
+        return torch.matmul(visual_tokens, self.weight)
+
+
+class Ovis2PreTrainedModel(PreTrainedModel):
+    config: Ovis2Config
+    base_model_prefix = "model"
+    input_modalities = ("image", "text")
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["Ovis2VisionAttention"]
+    _skip_keys_device_placement = ["past_key_values"]
+    _supports_cache_class = True
+    _supports_flash_attn = True
+    _supports_flex_attn = True
+    _supports_sdpa = True
+
+    _can_compile_fullgraph = True
+    _supports_attention_backend = True
+
+    def _init_weights(self, module):
+        super()._init_weights(module)
+        if isinstance(module, Ovis2VisionEmbeddings):
+            init.copy_(module.position_ids, torch.arange(module.position_ids.shape[-1]).expand((1, -1)))
+
+
+def hard_softmax(logits: torch.Tensor, dim: int):
+    y_soft = logits.softmax(dim)
+    # Straight through.
+    index = y_soft.max(dim, keepdim=True)[1]
+    y_hard = torch.zeros_like(logits, memory_format=torch.legacy_contiguous_format).scatter_(dim, index, 1.0)
+    ret = y_hard - y_soft.detach() + y_soft
+
+    return ret
+
+
+class Ovis2VisionModel(Ovis2PreTrainedModel):
+    config: Ovis2VisionConfig
+    _can_record_outputs = {
+        "hidden_states": Ovis2VisionEncoderLayer,
+        "attentions": Ovis2VisionAttention,
+    }
+
+    def __init__(self, config: Ovis2VisionConfig):
+        super().__init__(config)
+        self.config = config
+        self.transformer = Ovis2VisionTransformer(config)
+        self.num_visual_indicator_tokens = config.num_visual_indicator_tokens
+        self.vocab_size = config.vocab_size
+        self.head_linear = nn.Linear(
+            config.hidden_size * config.hidden_stride * config.hidden_stride,
+            self.vocab_size - self.num_visual_indicator_tokens,
+            bias=False,
+        )
+        self.head_norm = nn.LayerNorm(self.vocab_size - self.num_visual_indicator_tokens)
+
+        self.post_init()
+
+    @merge_with_config_defaults
+    @capture_outputs
+    def forward(
+        self, pixel_values: torch.FloatTensor, **kwargs: Unpack[TransformersKwargs]
+    ) -> tuple | BaseModelOutputWithVisualIndicatorFeatures:
+        outputs = self.transformer(pixel_values, **kwargs)
+        last_hidden_state = outputs[0]
+        if self.config.hidden_stride > 1:
+            num_images, seq_len, hidden_dim = last_hidden_state.shape
+            hidden_stride = self.config.hidden_stride
+
+            sqrt_l = int(math.sqrt(seq_len))
+            if sqrt_l * sqrt_l != seq_len:
+                raise ValueError("Token sequence length must be a perfect square")
+
+            pad_size = (hidden_stride - (sqrt_l % hidden_stride)) % hidden_stride
+            last_hidden_state = nn.functional.pad(last_hidden_state, (0, 0, 0, pad_size, 0, pad_size), "constant", 0)
+            sqrt_l += pad_size
+
+            last_hidden_state = last_hidden_state.reshape(
+                num_images, sqrt_l // hidden_stride, hidden_stride, sqrt_l // hidden_stride, hidden_stride, hidden_dim
+            )
+            last_hidden_state = last_hidden_state.permute(0, 1, 3, 2, 4, 5)
+            last_hidden_state = last_hidden_state.reshape(
+                num_images, -1, hidden_stride * hidden_stride * hidden_dim
+            )  # (n, (sqrt_l//hs)^2, hs^2*d)
+
+        logits = self.head_linear(last_hidden_state)
+        logits = self.head_norm(logits)
+
+        if self.config.tokenize_function == "gumbel_argmax":
+            prob_token = nn.functional.gumbel_softmax(logits, dim=-1, hard=True)
+        elif self.config.tokenize_function == "st_argmax":
+            prob_token = hard_softmax(logits, dim=-1)
+        elif self.config.tokenize_function == "softmax":
+            prob_token = nn.functional.softmax(logits, dim=-1)
+
+        return BaseModelOutputWithVisualIndicatorFeatures(
+            last_hidden_state=last_hidden_state,
+            pooler_output=prob_token,
+        )
+
+
+@auto_docstring(
+    custom_intro="""
+    The Ovis2 model which consists of a vision backbone and a language model, without a language modeling head.
+    """
+)
+class Ovis2Model(Ovis2PreTrainedModel):
+    def __init__(self, config: Ovis2Config):
+        super().__init__(config)
+        self.vision_tower = Ovis2VisionModel(config.vision_config)
+        self.language_model = AutoModel.from_config(config.text_config)
+        self.visual_embeddings_table = Ovis2VisualEmbeddingTable(config.vision_config.vocab_size, config.hidden_size)
+
+        self.visual_vocab_size = config.vision_config.vocab_size
+        self.vocab_size = config.vocab_size
+        self.visual_indicator_token_ids = config.visual_indicator_token_ids
+        self.post_init()
+
+    @can_return_tuple
+    @auto_docstring(
+        custom_intro="Obtains image last hidden states from the vision tower and apply multimodal projection."
+    )
+    def get_image_features(
+        self,
+        pixel_values: torch.FloatTensor,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> tuple | BaseModelOutputWithVisualIndicatorFeatures:
+        image_outputs = self.vision_tower(pixel_values, return_dict=True, **kwargs)
+        image_features = image_outputs.pooler_output
+        batch_size, img_seq_len, _ = image_features.shape
+        padding_tensor = torch.zeros(
+            (batch_size, img_seq_len, self.vision_tower.num_visual_indicator_tokens),
+            dtype=image_features.dtype,
+            device=image_features.device,
+            requires_grad=False,
+            layout=image_features.layout,
+        )
+        image_features = torch.cat([image_features, padding_tensor], dim=2)
+        image_features = self.visual_embeddings_table(image_features)
+
+        visual_indicator = torch.arange(
+            self.visual_vocab_size - self.vision_tower.num_visual_indicator_tokens,
+            self.visual_vocab_size,
+            dtype=torch.long,
+        ).to(image_features.device)
+        image_outputs.pooler_output = image_features
+        image_outputs.visual_indicator_features = self.visual_embeddings_table(visual_indicator)
+
+        return image_outputs
+
+    def get_placeholder_mask(
+        self, input_ids: torch.LongTensor, inputs_embeds: torch.FloatTensor, image_features: torch.FloatTensor
+    ):
+        """
+        Obtains multimodal placeholder mask from `input_ids` or `inputs_embeds`, and checks that the placeholder token count is
+        equal to the length of multimodal features. If the lengths are different, an error is raised.
+        """
+        if input_ids is None:
+            special_image_mask = inputs_embeds == self.get_input_embeddings()(
+                torch.tensor(self.config.image_token_id, dtype=torch.long, device=inputs_embeds.device)
+            )
+            special_image_mask = special_image_mask.all(-1)
+        else:
+            special_image_mask = input_ids == self.config.image_token_id
+
+        n_image_tokens = special_image_mask.sum()
+        n_image_features = image_features.shape[0] * image_features.shape[1]
+        special_image_mask = special_image_mask.unsqueeze(-1).expand_as(inputs_embeds).to(inputs_embeds.device)
+        torch_compilable_check(
+            inputs_embeds[special_image_mask].numel() == image_features.numel(),
+            f"Image features and image tokens do not match, tokens: {n_image_tokens}, features: {n_image_features}",
+        )
+        return special_image_mask
+
+    @can_return_tuple
+    @auto_docstring
+    def forward(
+        self,
+        input_ids: torch.LongTensor | None = None,
+        pixel_values: torch.FloatTensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values: Cache | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        labels: torch.LongTensor | None = None,
+        use_cache: bool | None = None,
+        logits_to_keep: int | torch.Tensor = 0,
+        **kwargs,
+    ) -> tuple | Ovis2ModelOutputWithPast:
+        if (input_ids is None) ^ (inputs_embeds is not None):
+            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+        if inputs_embeds is None:
+            inputs_embeds = self.get_input_embeddings()(input_ids)
+
+        if pixel_values is not None:
+            image_outputs = self.get_image_features(pixel_values=pixel_values, return_dict=True)
+            image_features = image_outputs.pooler_output
+            visual_indicator_features = image_outputs.visual_indicator_features
+
+            special_image_mask = self.get_placeholder_mask(
+                input_ids,
+                inputs_embeds=inputs_embeds,
+                image_features=image_features,
+            )
+            inputs_embeds = inputs_embeds.masked_scatter(special_image_mask, image_features)
+
+            for i, visual_indicator_id in enumerate(self.visual_indicator_token_ids):
+                if input_ids is None:
+                    mask = inputs_embeds == self.get_input_embeddings()(
+                        torch.tensor(visual_indicator_id, dtype=torch.long, device=inputs_embeds.device)
+                    )
+                    mask = mask.all(-1)
+                else:
+                    mask = (input_ids == visual_indicator_id).to(inputs_embeds.device)
+
+                if mask.any():
+                    inputs_embeds[mask] = (
+                        visual_indicator_features[i]
+                        .expand_as(inputs_embeds[mask])
+                        .to(inputs_embeds.device, inputs_embeds.dtype)
+                    )
+
+        outputs = self.language_model(
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            logits_to_keep=logits_to_keep,
+            **kwargs,
+        )
+
+        return Ovis2ModelOutputWithPast(
+            last_hidden_state=outputs.last_hidden_state,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+            image_hidden_states=image_features if pixel_values is not None else None,
+        )
+
+
+@auto_docstring
+class Ovis2ForConditionalGeneration(Ovis2PreTrainedModel, GenerationMixin):
+    _tied_weights_keys = {"lm_head.weight": "model.language_model.embed_tokens.weight"}
+
+    def __init__(self, config: Ovis2Config):
+        super().__init__(config)
+        self.model = Ovis2Model(config)
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.post_init()
+
+    def get_output_embeddings(self) -> nn.Module:
+        return self.lm_head
+
+    @auto_docstring
+    def get_image_features(
+        self, pixel_values: torch.FloatTensor, **kwargs: Unpack[TransformersKwargs]
+    ) -> tuple | BaseModelOutputWithVisualIndicatorFeatures:
+        return self.model.get_image_features(pixel_values=pixel_values, **kwargs)
+
+    @can_return_tuple
+    @auto_docstring
+    def forward(
+        self,
+        input_ids: torch.LongTensor | None = None,
+        pixel_values: torch.FloatTensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values: Cache | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        labels: torch.LongTensor | None = None,
+        use_cache: bool | None = None,
+        logits_to_keep: int | torch.Tensor = 0,
+        **kwargs,
+    ) -> tuple | Ovis2CausalLMOutputWithPast:
+        r"""
+        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
+            config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
+            (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
+
+        Example:
+
+        ```python
+        >>> from PIL import Image
+        >>> import httpx
+        >>> from io import BytesIO
+        >>> from transformers import AutoProcessor, Ovis2ForConditionalGeneration
+
+        >>> model = Ovis2ForConditionalGeneration.from_pretrained("thisisiron/Ovis2-2B-hf")
+        >>> processor = AutoProcessor.from_pretrained("thisisiron/Ovis2-2B-hf")
+
+        >>> prompt = "<|im_start|>user\n<image>\nDescribe the image.<|im_end|>\n<|im_start|>assistant\n"
+        >>> url = "http://images.cocodataset.org/val2014/COCO_val2014_000000537955.jpg"
+        >>> with httpx.stream("GET", url) as response:
+        ...     image = Image.open(BytesIO(response.read()))
+
+        >>> inputs = processor(images=image, text=prompt, return_tensors="pt")
+
+        >>> # Generate
+        >>> generate_ids = model.generate(**inputs, max_new_tokens=15)
+        >>> processor.batch_decode(generate_ids, skip_special_tokens=True)[0]
+        "user\n\nDescribe the image.\nassistant\nThe image features a brown dog standing on a wooden floor, looking up with"
+        ```"""
+        outputs = self.model(
+            input_ids=input_ids,
+            pixel_values=pixel_values,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            **kwargs,
+        )
+
+        hidden_states = outputs[0]
+        # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
+        slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+        logits = self.lm_head(hidden_states[:, slice_indices, :])
+
+        loss = None
+        if labels is not None:
+            loss = self.loss_function(
+                logits=logits, labels=labels, vocab_size=self.config.text_config.vocab_size, **kwargs
+            )
+
+        return Ovis2CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+            image_hidden_states=outputs.image_hidden_states,
+        )
+
+    def prepare_inputs_for_generation(
+        self,
+        input_ids,
+        past_key_values=None,
+        inputs_embeds=None,
+        pixel_values=None,
+        attention_mask=None,
+        logits_to_keep=None,
+        is_first_iteration=False,
+        **kwargs,
+    ):
+        # Overwritten -- in specific circumstances we don't want to forward image inputs to the model
+
+        model_inputs = super().prepare_inputs_for_generation(
+            input_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            attention_mask=attention_mask,
+            logits_to_keep=logits_to_keep,
+            is_first_iteration=is_first_iteration,
+            **kwargs,
+        )
+
+        if is_first_iteration or not kwargs.get("use_cache", True):
+            # Pixel values are used only in the first iteration if available
+            # In subsequent iterations, they are already merged with text and cached
+            # NOTE: first iteration doesn't have to be prefill, it can be the first
+            # iteration with a question and cached system prompt (continue generate from cache)
+            model_inputs["pixel_values"] = pixel_values
+
+        return model_inputs
+
+
+__all__ = ["Ovis2PreTrainedModel", "Ovis2Model", "Ovis2ForConditionalGeneration"]

LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/textnet/__init__.py

@@ -0,0 +1,29 @@
+# Copyright 2024 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from typing import TYPE_CHECKING
+
+from ...utils import _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+    from .configuration_textnet import *
+    from .image_processing_pil_textnet import *
+    from .image_processing_textnet import *
+    from .modeling_textnet import *
+else:
+    import sys
+
+    _file = globals()["__file__"]
+    sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)

LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/textnet/configuration_textnet.py

@@ -0,0 +1,94 @@
+# Copyright 2024 the Fast authors and HuggingFace Inc. team.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""TextNet model configuration"""
+
+from huggingface_hub.dataclasses import strict
+
+from ...backbone_utils import BackboneConfigMixin
+from ...configuration_utils import PreTrainedConfig
+from ...utils import auto_docstring
+
+
+@auto_docstring(checkpoint="czczup/textnet-base")
+@strict
+class TextNetConfig(BackboneConfigMixin, PreTrainedConfig):
+    r"""
+    stem_kernel_size (`int`, *optional*, defaults to 3):
+        The kernel size for the initial convolution layer.
+    stem_stride (`int`, *optional*, defaults to 2):
+        The stride for the initial convolution layer.
+    stem_num_channels (`int`, *optional*, defaults to 3):
+        The num of channels in input for the initial convolution layer.
+    stem_out_channels (`int`, *optional*, defaults to 64):
+        The num of channels in out for the initial convolution layer.
+    stem_act_func (`str`, *optional*, defaults to `"relu"`):
+        The activation function for the initial convolution layer.
+    conv_layer_kernel_sizes (`list[list[list[int]]]`, *optional*):
+        A list of stage-wise kernel sizes. If `None`, defaults to:
+        `[[[3, 3], [3, 3], [3, 3]], [[3, 3], [1, 3], [3, 3], [3, 1]], [[3, 3], [3, 3], [3, 1], [1, 3]], [[3, 3], [3, 1], [1, 3], [3, 3]]]`.
+    conv_layer_strides (`list[list[int]]`, *optional*):
+        A list of stage-wise strides. If `None`, defaults to:
+        `[[1, 2, 1], [2, 1, 1, 1], [2, 1, 1, 1], [2, 1, 1, 1]]`.
+
+    Examples:
+
+    ```python
+    >>> from transformers import TextNetConfig, TextNetBackbone
+
+    >>> # Initializing a TextNetConfig
+    >>> configuration = TextNetConfig()
+
+    >>> # Initializing a model (with random weights)
+    >>> model = TextNetBackbone(configuration)
+
+    >>> # Accessing the model configuration
+    >>> configuration = model.config
+    ```"""
+
+    model_type = "textnet"
+
+    stem_kernel_size: int = 3
+    stem_stride: int = 2
+    stem_num_channels: int = 3
+    stem_out_channels: int = 64
+    stem_act_func: str = "relu"
+    image_size: list[int] | tuple[int, int] | int = (640, 640)
+    conv_layer_kernel_sizes: list | None = None
+    conv_layer_strides: list | None = None
+    hidden_sizes: list[int] | tuple[int, ...] = (64, 64, 128, 256, 512)
+    batch_norm_eps: float = 1e-5
+    initializer_range: float = 0.02
+    _out_features: list[str] | None = None
+    _out_indices: list[int] | None = None
+
+    def __post_init__(self, **kwargs):
+        if self.conv_layer_kernel_sizes is None:
+            self.conv_layer_kernel_sizes = [
+                [[3, 3], [3, 3], [3, 3]],
+                [[3, 3], [1, 3], [3, 3], [3, 1]],
+                [[3, 3], [3, 3], [3, 1], [1, 3]],
+                [[3, 3], [3, 1], [1, 3], [3, 3]],
+            ]
+        if self.conv_layer_strides is None:
+            self.conv_layer_strides = [[1, 2, 1], [2, 1, 1, 1], [2, 1, 1, 1], [2, 1, 1, 1]]
+
+        self.depths = [len(layer) for layer in self.conv_layer_kernel_sizes]
+        self.stage_names = ["stem"] + [f"stage{idx}" for idx in range(1, 5)]
+        self.set_output_features_output_indices(
+            out_indices=kwargs.pop("out_indices", None), out_features=kwargs.pop("out_features", None)
+        )
+        super().__post_init__(**kwargs)
+
+
+__all__ = ["TextNetConfig"]

LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/textnet/image_processing_pil_textnet.py

@@ -0,0 +1,134 @@
+# Copyright 2024 the Fast authors and The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Image processor class for TextNet."""
+
+import numpy as np
+
+from ...image_processing_backends import PilBackend
+from ...image_processing_utils import BatchFeature
+from ...image_transforms import get_resize_output_image_size
+from ...image_transforms import resize as np_resize
+from ...image_utils import (
+    IMAGENET_DEFAULT_MEAN,
+    IMAGENET_DEFAULT_STD,
+    ChannelDimension,
+    ImageInput,
+    PILImageResampling,
+    SizeDict,
+)
+from ...processing_utils import ImagesKwargs, Unpack
+from ...utils import TensorType, auto_docstring
+
+
+# Adapted from transformers.models.textnet.image_processing_textnet.TextNetImageProcessorKwargs
+class TextNetImageProcessorKwargs(ImagesKwargs, total=False):
+    """
+    size_divisor (`int`, *optional*, defaults to `self.size_divisor`):
+        Ensures height and width are rounded to a multiple of this value after resizing.
+    """
+
+    size_divisor: int
+
+
+@auto_docstring
+class TextNetImageProcessorPil(PilBackend):
+    """PIL backend for TextNet with size_divisor resize."""
+
+    valid_kwargs = TextNetImageProcessorKwargs
+
+    resample = PILImageResampling.BILINEAR
+    image_mean = IMAGENET_DEFAULT_MEAN
+    image_std = IMAGENET_DEFAULT_STD
+    size = {"shortest_edge": 640}
+    default_to_square = False
+    crop_size = {"height": 224, "width": 224}
+    do_resize = True
+    do_center_crop = False
+    do_rescale = True
+    do_normalize = True
+    do_convert_rgb = True
+    size_divisor = 32
+
+    def __init__(self, **kwargs: Unpack[TextNetImageProcessorKwargs]):
+        super().__init__(**kwargs)
+
+    @auto_docstring
+    def preprocess(self, images: ImageInput, **kwargs: Unpack[TextNetImageProcessorKwargs]) -> BatchFeature:
+        return super().preprocess(images, **kwargs)
+
+    def resize(
+        self,
+        image: np.ndarray,
+        size: SizeDict,
+        resample: "PILImageResampling | None",
+        size_divisor: int = 32,
+        **kwargs,
+    ) -> np.ndarray:
+        """Resize to shortest_edge then round up to be divisible by size_divisor."""
+        if not size.shortest_edge:
+            raise ValueError(f"Size must contain 'shortest_edge' key. Got {size.keys()}")
+        height, width = get_resize_output_image_size(
+            image,
+            size=size.shortest_edge,
+            default_to_square=False,
+            input_data_format=ChannelDimension.FIRST,
+        )
+        # Round up to be divisible by size_divisor
+        if height % size_divisor != 0:
+            height += size_divisor - (height % size_divisor)
+        if width % size_divisor != 0:
+            width += size_divisor - (width % size_divisor)
+        return np_resize(
+            image,
+            size=(height, width),
+            resample=resample,
+            data_format=ChannelDimension.FIRST,
+            input_data_format=ChannelDimension.FIRST,
+        )
+
+    def _preprocess(
+        self,
+        images: list[np.ndarray],
+        do_resize: bool,
+        size: SizeDict,
+        resample: "PILImageResampling | None",
+        do_center_crop: bool,
+        crop_size: SizeDict,
+        do_rescale: bool,
+        rescale_factor: float,
+        do_normalize: bool,
+        image_mean: float | list[float] | None,
+        image_std: float | list[float] | None,
+        do_pad: bool | None,
+        pad_size: SizeDict | None,
+        return_tensors: str | TensorType | None,
+        size_divisor: int = 32,
+        **kwargs,
+    ) -> BatchFeature:
+        """Custom preprocessing for TextNet."""
+        processed_images = []
+        for image in images:
+            if do_resize:
+                image = self.resize(image, size, resample, size_divisor=size_divisor)
+            if do_center_crop:
+                image = self.center_crop(image, crop_size)
+            if do_rescale:
+                image = self.rescale(image, rescale_factor)
+            if do_normalize:
+                image = self.normalize(image, image_mean, image_std)
+            processed_images.append(image)
+        return BatchFeature(data={"pixel_values": processed_images}, tensor_type=return_tensors)
+
+
+__all__ = ["TextNetImageProcessorPil"]

LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/textnet/modeling_textnet.py

@@ -0,0 +1,417 @@
+# Copyright 2024 the Fast authors and The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""PyTorch TextNet model."""
+
+from typing import Any
+
+import torch
+import torch.nn as nn
+from torch import Tensor
+
+from ...activations import ACT2CLS
+from ...backbone_utils import BackboneMixin, filter_output_hidden_states
+from ...modeling_outputs import (
+    BackboneOutput,
+    BaseModelOutputWithNoAttention,
+    BaseModelOutputWithPoolingAndNoAttention,
+    ImageClassifierOutputWithNoAttention,
+)
+from ...modeling_utils import PreTrainedModel
+from ...utils import auto_docstring, logging
+from ...utils.generic import can_return_tuple
+from .configuration_textnet import TextNetConfig
+
+
+logger = logging.get_logger(__name__)
+
+
+class TextNetConvLayer(nn.Module):
+    def __init__(self, config: TextNetConfig):
+        super().__init__()
+
+        self.kernel_size = config.stem_kernel_size
+        self.stride = config.stem_stride
+        self.activation_function = config.stem_act_func
+
+        padding = (
+            (config.kernel_size[0] // 2, config.kernel_size[1] // 2)
+            if isinstance(config.stem_kernel_size, tuple)
+            else config.stem_kernel_size // 2
+        )
+
+        self.conv = nn.Conv2d(
+            config.stem_num_channels,
+            config.stem_out_channels,
+            kernel_size=config.stem_kernel_size,
+            stride=config.stem_stride,
+            padding=padding,
+            bias=False,
+        )
+        self.batch_norm = nn.BatchNorm2d(config.stem_out_channels, config.batch_norm_eps)
+
+        self.activation = nn.Identity()
+        if self.activation_function is not None:
+            self.activation = ACT2CLS[self.activation_function]()
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.conv(hidden_states)
+        hidden_states = self.batch_norm(hidden_states)
+        return self.activation(hidden_states)
+
+
+class TextNetRepConvLayer(nn.Module):
+    r"""
+    This layer supports re-parameterization by combining multiple convolutional branches
+    (e.g., main convolution, vertical, horizontal, and identity branches) during training.
+    At inference time, these branches can be collapsed into a single convolution for
+    efficiency, as per the re-parameterization paradigm.
+
+    The "Rep" in the name stands for "re-parameterization" (introduced by RepVGG).
+    """
+
+    def __init__(self, config: TextNetConfig, in_channels: int, out_channels: int, kernel_size: int, stride: int):
+        super().__init__()
+
+        self.num_channels = in_channels
+        self.out_channels = out_channels
+        self.kernel_size = kernel_size
+        self.stride = stride
+
+        padding = ((kernel_size[0] - 1) // 2, (kernel_size[1] - 1) // 2)
+
+        self.activation_function = nn.ReLU()
+
+        self.main_conv = nn.Conv2d(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            kernel_size=kernel_size,
+            stride=stride,
+            padding=padding,
+            bias=False,
+        )
+        self.main_batch_norm = nn.BatchNorm2d(num_features=out_channels, eps=config.batch_norm_eps)
+
+        vertical_padding = ((kernel_size[0] - 1) // 2, 0)
+        horizontal_padding = (0, (kernel_size[1] - 1) // 2)
+
+        if kernel_size[1] != 1:
+            self.vertical_conv = nn.Conv2d(
+                in_channels=in_channels,
+                out_channels=out_channels,
+                kernel_size=(kernel_size[0], 1),
+                stride=stride,
+                padding=vertical_padding,
+                bias=False,
+            )
+            self.vertical_batch_norm = nn.BatchNorm2d(num_features=out_channels, eps=config.batch_norm_eps)
+        else:
+            self.vertical_conv, self.vertical_batch_norm = None, None
+
+        if kernel_size[0] != 1:
+            self.horizontal_conv = nn.Conv2d(
+                in_channels=in_channels,
+                out_channels=out_channels,
+                kernel_size=(1, kernel_size[1]),
+                stride=stride,
+                padding=horizontal_padding,
+                bias=False,
+            )
+            self.horizontal_batch_norm = nn.BatchNorm2d(num_features=out_channels, eps=config.batch_norm_eps)
+        else:
+            self.horizontal_conv, self.horizontal_batch_norm = None, None
+
+        self.rbr_identity = (
+            nn.BatchNorm2d(num_features=in_channels, eps=config.batch_norm_eps)
+            if out_channels == in_channels and stride == 1
+            else None
+        )
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        main_outputs = self.main_conv(hidden_states)
+        main_outputs = self.main_batch_norm(main_outputs)
+
+        # applies a convolution with a vertical kernel
+        if self.vertical_conv is not None:
+            vertical_outputs = self.vertical_conv(hidden_states)
+            vertical_outputs = self.vertical_batch_norm(vertical_outputs)
+            main_outputs = main_outputs + vertical_outputs
+
+        # applies a convolution with a horizontal kernel
+        if self.horizontal_conv is not None:
+            horizontal_outputs = self.horizontal_conv(hidden_states)
+            horizontal_outputs = self.horizontal_batch_norm(horizontal_outputs)
+            main_outputs = main_outputs + horizontal_outputs
+
+        if self.rbr_identity is not None:
+            id_out = self.rbr_identity(hidden_states)
+            main_outputs = main_outputs + id_out
+
+        return self.activation_function(main_outputs)
+
+
+class TextNetStage(nn.Module):
+    def __init__(self, config: TextNetConfig, depth: int):
+        super().__init__()
+        kernel_size = config.conv_layer_kernel_sizes[depth]
+        stride = config.conv_layer_strides[depth]
+
+        num_layers = len(kernel_size)
+        stage_in_channel_size = config.hidden_sizes[depth]
+        stage_out_channel_size = config.hidden_sizes[depth + 1]
+
+        in_channels = [stage_in_channel_size] + [stage_out_channel_size] * (num_layers - 1)
+        out_channels = [stage_out_channel_size] * num_layers
+
+        stage = []
+        for stage_config in zip(in_channels, out_channels, kernel_size, stride):
+            stage.append(TextNetRepConvLayer(config, *stage_config))
+        self.stage = nn.ModuleList(stage)
+
+    def forward(self, hidden_state):
+        for block in self.stage:
+            hidden_state = block(hidden_state)
+        return hidden_state
+
+
+class TextNetEncoder(nn.Module):
+    def __init__(self, config: TextNetConfig):
+        super().__init__()
+
+        stages = []
+        num_stages = len(config.conv_layer_kernel_sizes)
+        for stage_ix in range(num_stages):
+            stages.append(TextNetStage(config, stage_ix))
+
+        self.stages = nn.ModuleList(stages)
+
+    def forward(
+        self,
+        hidden_state: torch.Tensor,
+        output_hidden_states: bool | None = None,
+        return_dict: bool | None = None,
+    ) -> BaseModelOutputWithNoAttention:
+        hidden_states = [hidden_state]
+        for stage in self.stages:
+            hidden_state = stage(hidden_state)
+            hidden_states.append(hidden_state)
+
+        if not return_dict:
+            output = (hidden_state,)
+            return output + (hidden_states,) if output_hidden_states else output
+
+        return BaseModelOutputWithNoAttention(last_hidden_state=hidden_state, hidden_states=hidden_states)
+
+
+@auto_docstring
+class TextNetPreTrainedModel(PreTrainedModel):
+    config: TextNetConfig
+    base_model_prefix = "textnet"
+    main_input_name = "pixel_values"
+
+
+@auto_docstring
+class TextNetModel(TextNetPreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+        self.stem = TextNetConvLayer(config)
+        self.encoder = TextNetEncoder(config)
+        self.pooler = nn.AdaptiveAvgPool2d((2, 2))
+        self.post_init()
+
+    @auto_docstring
+    def forward(
+        self,
+        pixel_values: Tensor,
+        output_hidden_states: bool | None = None,
+        return_dict: bool | None = None,
+        **kwargs,
+    ) -> tuple[Any, list[Any]] | tuple[Any] | BaseModelOutputWithPoolingAndNoAttention:
+        return_dict = return_dict if return_dict is not None else self.config.return_dict
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+
+        hidden_state = self.stem(pixel_values)
+
+        encoder_outputs = self.encoder(
+            hidden_state, output_hidden_states=output_hidden_states, return_dict=return_dict
+        )
+
+        last_hidden_state = encoder_outputs[0]
+        pooled_output = self.pooler(last_hidden_state)
+
+        if not return_dict:
+            output = (last_hidden_state, pooled_output)
+            return output + (encoder_outputs[1],) if output_hidden_states else output
+
+        return BaseModelOutputWithPoolingAndNoAttention(
+            last_hidden_state=last_hidden_state,
+            pooler_output=pooled_output,
+            hidden_states=encoder_outputs[1] if output_hidden_states else None,
+        )
+
+
+@auto_docstring(
+    custom_intro="""
+    TextNet Model with an image classification head on top (a linear layer on top of the pooled features), e.g. for
+    ImageNet.
+    """
+)
+class TextNetForImageClassification(TextNetPreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+        self.textnet = TextNetModel(config)
+        self.avg_pool = nn.AdaptiveAvgPool2d((1, 1))
+        self.flatten = nn.Flatten()
+        self.fc = nn.Linear(config.hidden_sizes[-1], config.num_labels) if config.num_labels > 0 else nn.Identity()
+
+        # classification head
+        self.classifier = nn.ModuleList([self.avg_pool, self.flatten])
+
+        # initialize weights and apply final processing
+        self.post_init()
+
+    @auto_docstring
+    def forward(
+        self,
+        pixel_values: torch.FloatTensor | None = None,
+        labels: torch.LongTensor | None = None,
+        output_hidden_states: bool | None = None,
+        return_dict: bool | None = None,
+        **kwargs,
+    ) -> ImageClassifierOutputWithNoAttention:
+        r"""
+        labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
+            Labels for computing the image classification/regression loss. Indices should be in `[0, ...,
+            config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
+            `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
+
+        Examples:
+        ```python
+        >>> import torch
+        >>> import httpx
+        >>> from io import BytesIO
+        >>> from transformers import TextNetForImageClassification, TextNetImageProcessor
+        >>> from PIL import Image
+
+        >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
+        >>> with httpx.stream("GET", url) as response:
+        ...     image = Image.open(BytesIO(response.read()))
+
+        >>> processor = TextNetImageProcessor.from_pretrained("czczup/textnet-base")
+        >>> model = TextNetForImageClassification.from_pretrained("czczup/textnet-base")
+
+        >>> inputs = processor(images=image, return_tensors="pt")
+        >>> with torch.no_grad():
+        ...     outputs = model(**inputs)
+        >>> outputs.logits.shape
+        torch.Size([1, 2])
+        ```"""
+        return_dict = return_dict if return_dict is not None else self.config.return_dict
+
+        outputs = self.textnet(pixel_values, output_hidden_states=output_hidden_states, return_dict=return_dict)
+        last_hidden_state = outputs[0]
+        for layer in self.classifier:
+            last_hidden_state = layer(last_hidden_state)
+        logits = self.fc(last_hidden_state)
+        loss = None
+
+        if labels is not None:
+            loss = self.loss_function(labels, logits, self.config)
+
+        if not return_dict:
+            output = (logits,) + outputs[2:]
+            return (loss,) + output if loss is not None else output
+
+        return ImageClassifierOutputWithNoAttention(loss=loss, logits=logits, hidden_states=outputs.hidden_states)
+
+
+@auto_docstring(
+    custom_intro="""
+    TextNet backbone, to be used with frameworks like DETR and MaskFormer.
+    """
+)
+class TextNetBackbone(BackboneMixin, TextNetPreTrainedModel):
+    has_attentions = False
+
+    def __init__(self, config):
+        super().__init__(config)
+
+        self.textnet = TextNetModel(config)
+        self.num_features = config.hidden_sizes
+
+        # initialize weights and apply final processing
+        self.post_init()
+
+    @can_return_tuple
+    @filter_output_hidden_states
+    @auto_docstring
+    def forward(
+        self,
+        pixel_values: Tensor,
+        output_hidden_states: bool | None = None,
+        return_dict: bool | None = None,
+        **kwargs,
+    ) -> tuple[tuple] | BackboneOutput:
+        r"""
+        Examples:
+
+        ```python
+        >>> import torch
+        >>> import httpx
+        >>> from io import BytesIO
+        >>> from PIL import Image
+        >>> from transformers import AutoImageProcessor, AutoBackbone
+
+        >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
+        >>> with httpx.stream("GET", url) as response:
+        ...     image = Image.open(BytesIO(response.read()))
+
+        >>> processor = AutoImageProcessor.from_pretrained("czczup/textnet-base")
+        >>> model = AutoBackbone.from_pretrained("czczup/textnet-base")
+
+        >>> inputs = processor(image, return_tensors="pt")
+        >>> with torch.no_grad():
+        >>>     outputs = model(**inputs)
+        ```"""
+        return_dict = return_dict if return_dict is not None else self.config.return_dict
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+
+        outputs = self.textnet(pixel_values, output_hidden_states=True, return_dict=return_dict)
+
+        hidden_states = outputs.hidden_states if return_dict else outputs[2]
+
+        feature_maps = ()
+        for idx, stage in enumerate(self.stage_names):
+            if stage in self.out_features:
+                feature_maps += (hidden_states[idx],)
+
+        if not return_dict:
+            output = (feature_maps,)
+            if output_hidden_states:
+                hidden_states = outputs.hidden_states if return_dict else outputs[2]
+                output += (hidden_states,)
+            return output
+
+        return BackboneOutput(
+            feature_maps=feature_maps,
+            hidden_states=outputs.hidden_states if output_hidden_states else None,
+            attentions=None,
+        )
+
+
+__all__ = ["TextNetBackbone", "TextNetModel", "TextNetPreTrainedModel", "TextNetForImageClassification"]