ptwm 1.0.0__cp313-cp313-macosx_11_0_arm64.whl

ptwm/__init__.py

@@ -0,0 +1,52 @@
+"""PTWM — lossless compression for PyTorch model weights.
+
+A typed preprocessing graph (bit reorder, byte/nibble split) feeds two
+peer entropy coders — Huffman and rANS — alongside a Zstd plane codec
+and a dedicated microscaling-format codec. The on-disk format is a
+multi-tensor ``.ptwm`` container with per-tensor random-access lookup.
+
+Public entry points: :class:`Compressor`, :class:`Decompressor`,
+:class:`CompressionConfig`, :class:`DecompressionConfig`,
+:class:`Method`, :class:`Format`. Third-party integrations
+(HuggingFace `transformers`, `safetensors`) live in :mod:`ptwm.integrations`
+and are imported from there explicitly, not re-exported here.
+"""
+
+from importlib.metadata import PackageNotFoundError, version
+
+from ._config import (
+    CompressionConfig,
+    DecompressionConfig,
+    Format,
+    Lossy,
+    Method,
+)
+from ._exceptions import (
+    CompressionMethodNotSupportedError,
+    Error,
+    HeaderParseError,
+    InvalidDTypeError,
+)
+from .core import Compressor, Decompressor
+from .integrations import materialize_hf_cache
+
+try:
+    __version__ = version("ptwm")
+except PackageNotFoundError:
+    __version__ = "0.0.0+unknown"
+
+__all__ = [
+    "CompressionConfig",
+    "CompressionMethodNotSupportedError",
+    "Compressor",
+    "DecompressionConfig",
+    "Decompressor",
+    "Format",
+    "Lossy",
+    "Method",
+    "Error",
+    "HeaderParseError",
+    "InvalidDTypeError",
+    "__version__",
+    "materialize_hf_cache",
+]

ptwm/_config.py

@@ -0,0 +1,162 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, replace
+from enum import Enum
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from ptwm._rust.policy import ResolvedPolicy
+
+    from .codecs import CodecId
+
+
+class Method(Enum):
+    AUTO = 0
+    HUFFMAN = 1
+    ZSTD = 2
+    IDENTITY = 3
+    RANS = 4
+    MICROSCALE = 5
+
+    @classmethod
+    def _missing_(cls, value):
+        if isinstance(value, str):
+            value = value.upper()
+            if value in cls.__members__:
+                return cls.__members__[value]
+        return None
+
+
+class Format(Enum):
+    BYTE = 1
+    TORCH = 2
+    NUMPY = 3
+    FILE = 4
+
+    @classmethod
+    def _missing_(cls, value):
+        if isinstance(value, str):
+            value = value.upper()
+            if value in cls.__members__:
+                return cls.__members__[value]
+        return None
+
+
+class Lossy(Enum):
+    NONE = 0
+    INTEGER = 1
+    UNSIGN = 2
+
+    @classmethod
+    def _missing_(cls, value):
+        if isinstance(value, str):
+            value = value.upper()
+            if value in cls.__members__:
+                return cls.__members__[value]
+        return None
+
+
+@dataclass(frozen=True, slots=True, kw_only=True)
+class CompressionConfig:
+    method: Method = Method.AUTO
+    input_format: Format = Format.BYTE
+    bytearray_dtype: str = "bfloat16"
+    is_monotonic: int = 0
+    threads: int | None = None
+    compression_threshold: float = 0.95
+    check_th_after_percent: int = 10
+    reorder_signbit: int = 0
+    delta_compressed_type: str | None = None
+    lossy_compressed_type: Lossy = Lossy.NONE
+    lossy_compressed_factor: int = 27
+    compression_chunk: int = 256 * 1024
+    is_streaming: bool = False
+    streaming_chunk: int = 1024 * 1024
+    input_file: str | None = None
+    compressed_file: str | None = None
+    decompressed_file: str | None = None
+    zstd_level: int = 3
+    lz4_compression_level: int = 0
+    codec_menu: list[CodecId] | None = None
+    """Restrict per-role trial-encode menus to this set of codecs.
+
+    None → use the full role-applicable menu. An empty intersection on a
+    role raises ``InvalidContainer`` — an Identity fallback would silently
+    falsify ablation measurements. Ignored when ``method`` selects a forced-
+    codec path (ZSTD / RANS / IDENTITY).
+    """
+
+    @classmethod
+    def from_resolved_policy(
+        cls,
+        rp: ResolvedPolicy,
+        *,
+        base: CompressionConfig | None = None,
+    ) -> CompressionConfig:
+        """Build a constrained ``CompressionConfig`` from a ``ResolvedPolicy``.
+
+        The returned config's ``codec_menu`` is the intersection of the
+        resolved policy's ``effective_global`` set with the wire-stable
+        built-in ``CodecId`` variants.
+
+        The resolved policy lists every contribution (built-in or installed
+        third-party) that is currently trusted *and* allowed by the policy
+        file. Restricting the trial-encode menu to that intersection is how
+        an ablation run actually exercises a variant policy: tensors that
+        previously round-tripped through a now-excluded codec will fall
+        through to the next-best allowed coder instead.
+
+        ``base`` is the starting config (defaults to a default-constructed
+        ``CompressionConfig``); every field other than ``codec_menu`` is
+        preserved.
+        """
+        # Imported locally — `ptwm.codecs` reaches back through
+        # `ptwm.preprocessing → ptwm.utils → ptwm._config` at top level,
+        # so a module-level import here is a real circular-import bug
+        # (not just a ruff style preference).
+        from ptwm._rust import builtin_canonical_id  # noqa: PLC0415
+
+        from .codecs import CodecId  # noqa: PLC0415
+
+        # Map every wire-stable built-in codec to the short name used in
+        # `crates/ptwm-core/src/extension/builtins.rs`.
+        _BUILTIN_CODEC_NAMES: dict[CodecId, str] = {
+            CodecId.Identity: "identity",
+            CodecId.Huffman: "huffman",
+            CodecId.Rans: "rans",
+            CodecId.Zstd: "zstd",
+            CodecId.PerGroupCodebook: "per_group_codebook",
+            CodecId.Order1ScaleAC: "order1_scale_ac",
+        }
+
+        allowed_hex = set(rp.effective_global_ids())
+
+        # Empty `effective_global` means "no policy applied" (the default
+        # `PolicyFile.default_empty()` case): preserve the base's
+        # codec_menu rather than collapsing to an empty list, which the
+        # Rust dispatcher would treat as "no codecs allowed → error".
+        if not allowed_hex:
+            return base if base is not None else cls()
+
+        menu: list[CodecId] = []
+        for codec_id, short_name in _BUILTIN_CODEC_NAMES.items():
+            if builtin_canonical_id(short_name).hex() in allowed_hex:
+                menu.append(codec_id)
+
+        starting = base if base is not None else cls()
+        return replace(starting, codec_menu=menu)
+
+
+@dataclass(frozen=True, slots=True, kw_only=True)
+class DecompressionConfig:
+    # Most properties are inferred from the header; this struct only carries
+    # runtime specifics like thread count and the delta reference.
+    threads: int | None = None
+    delta_second_data: bytes | None = None
+    skip_missing: bool = False
+    """If ``True``, open containers that reference non-builtin extensions
+    that are not installed on this host.  Tensors needing those extensions
+    will fail when decoded, but the container can still be opened and
+    interrogated (e.g. listing tensor names).  When ``False`` (the default),
+    opening such a container raises ``ValueError``.
+    """

ptwm/_exceptions.py

@@ -0,0 +1,26 @@
+import logging
+
+# Standard logger for the package
+logger = logging.getLogger("ptwm")
+# Default handler setup could be left to the user, but we can provide a NullHandler
+logger.addHandler(logging.NullHandler())
+
+
+class Error(Exception):
+    """Base exception for all errors originating from the PTWM library."""
+
+
+class CompressionMethodNotSupportedError(Error):
+    """Raised when an unsupported compression method is requested."""
+
+
+class InvalidDTypeError(Error):
+    """Raised when a tensor or numpy array has an unsupported dtype."""
+
+
+class HeaderParseError(Error):
+    """Raised when the binary header is invalid or corrupt."""
+
+
+class UnsupportedQuantConfigError(ValueError):
+    """Raised when ``hf_quant_config.json`` carries an unsupported algo."""

ptwm/_loader.py

@@ -0,0 +1,48 @@
+"""Discover Python-resident extensions via entry points.
+
+Entry points under the `ptwm.extensions` group are loaded at PTWM
+startup. Each one is a callable that, when invoked, calls
+ContributionRegistry.register(...) for each contribution it provides.
+"""
+
+from __future__ import annotations
+
+import sys
+from collections.abc import Iterable
+from importlib.metadata import EntryPoint, entry_points
+
+
+def discover_host_extensions() -> Iterable[tuple[str, EntryPoint]]:
+    """Yield (entry_point_name, ep) for each ptwm.extensions entry point."""
+    eps = entry_points(group="ptwm.extensions")
+    for ep in eps:
+        yield ep.name, ep
+
+
+def load_all_host_extensions() -> int:
+    """Invoke every discovered entry point.
+
+    Returns the number of entry points successfully loaded. Errors from
+    a single entry point are logged (via print) but don't stop the
+    loader — a broken extension shouldn't take down the whole runtime.
+    """
+    count = 0
+    for name, ep in discover_host_extensions():
+        try:
+            register_fn = ep.load()
+        except Exception as exc:  # noqa: BLE001
+            print(  # noqa: T201
+                f"warning: ptwm.extensions[{name}] failed to load: {exc}",
+                file=sys.stderr,
+            )
+            continue
+        try:
+            register_fn()
+        except Exception as exc:  # noqa: BLE001
+            print(  # noqa: T201
+                f"warning: ptwm.extensions[{name}] failed to register: {exc}",
+                file=sys.stderr,
+            )
+            continue
+        count += 1
+    return count

ptwm/_rust/__init__.py

@@ -0,0 +1,313 @@
+"""Thin wrapper around the ``weights._core`` PyO3 extension.
+
+The ``weights._rust`` package is the ONLY place in the codebase that imports
+``weights._core``. Every other module goes through the wrappers defined here
+(and its submodules), so the underlying binding can swap (e.g. for a
+subprocess backend) without touching the rest of the codebase.
+"""
+
+from __future__ import annotations
+
+from ptwm import _core as _ext
+
+from ._errors import translate_errors
+
+__all__ = [
+    "blake3_hash",
+    "builtin_canonical_id",
+    "codec_id_registry",
+    "compress_model",
+    "decode_model",
+    "decode_tensor",
+    "decode_tensor_delta_hash",
+    "decode_tensor_info",
+    "decode_tensor_shape",
+    "delta_float_sub_decode",
+    "delta_float_sub_encode",
+    "delta_xor_decode",
+    "delta_xor_encode",
+    "dtype_element_size",
+    "explode_ptwm",
+    "huffman_decode",
+    "huffman_encode",
+    "identity_decode",
+    "identity_encode",
+    "implode_ptwm",
+    "list_plane_summaries",
+    "list_tensor_names",
+    "rans_decode",
+    "rans_encode",
+    "shannon",
+    "translate_errors",
+    "xxhash64",
+    "zstd_decode",
+    "zstd_encode",
+]
+
+
+@translate_errors
+def xxhash64(data: bytes | memoryview) -> int:
+    """Return the 64-bit xxhash of ``data`` (seed 0)."""
+    return _ext.xxhash64(data)
+
+
+@translate_errors
+def shannon(data: bytes | memoryview) -> float:
+    """Return the Shannon entropy of ``data`` in bits/byte."""
+    return _ext.shannon(data)
+
+
+# --- Per-plane codec encode / decode ---------------------------------------
+
+
+@translate_errors
+def identity_encode(data: bytes | memoryview) -> bytes:
+    return bytes(_ext.identity_encode(data))
+
+
+@translate_errors
+def identity_decode(blob: bytes | memoryview, expected_len: int) -> bytes:
+    return bytes(_ext.identity_decode(blob, expected_len))
+
+
+@translate_errors
+def huffman_encode(data: bytes | memoryview) -> bytes:
+    return bytes(_ext.huffman_encode(data))
+
+
+@translate_errors
+def huffman_decode(blob: bytes | memoryview, expected_len: int) -> bytes:
+    return bytes(_ext.huffman_decode(blob, expected_len))
+
+
+@translate_errors
+def rans_encode(data: bytes | memoryview) -> bytes:
+    return bytes(_ext.rans_encode(data))
+
+
+@translate_errors
+def rans_decode(blob: bytes | memoryview, expected_len: int) -> bytes:
+    return bytes(_ext.rans_decode(blob, expected_len))
+
+
+@translate_errors
+def zstd_encode(data: bytes | memoryview, level: int = 3) -> bytes:
+    return bytes(_ext.zstd_encode(data, level))
+
+
+@translate_errors
+def zstd_decode(blob: bytes | memoryview) -> bytes:
+    return bytes(_ext.zstd_decode(blob))
+
+
+# --- Delta primitives ------------------------------------------------------
+
+
+@translate_errors
+def delta_xor_encode(raw: bytes | memoryview, reference: bytes | memoryview) -> bytes:
+    return bytes(_ext.delta_xor_encode(raw, reference))
+
+
+@translate_errors
+def delta_xor_decode(
+    residual: bytes | memoryview, reference: bytes | memoryview
+) -> bytes:
+    return bytes(_ext.delta_xor_decode(residual, reference))
+
+
+@translate_errors
+def delta_float_sub_encode(
+    raw: bytes | memoryview, reference: bytes | memoryview
+) -> bytes:
+    return bytes(_ext.delta_float_sub_encode(raw, reference))
+
+
+@translate_errors
+def delta_float_sub_decode(
+    residual: bytes | memoryview, reference: bytes | memoryview
+) -> bytes:
+    return bytes(_ext.delta_float_sub_decode(residual, reference))
+
+
+@translate_errors
+def blake3_hash(data: bytes | memoryview) -> bytes:
+    """BLAKE3 digest (32 bytes) of ``data``."""
+    return bytes(_ext.blake3_hash(data))
+
+
+def codec_id_registry() -> list[tuple[str, int]]:
+    """Authoritative ``(name, wire_id)`` list from the Rust ``CodecId``.
+
+    The regression test that locks the Python ``CodecId`` IntEnum to the
+    Rust truth consumes this list.
+    """
+    return [(str(n), int(v)) for n, v in _ext.codec_id_registry()]
+
+
+def builtin_canonical_id(name: str) -> bytes:
+    r"""Return the 32-byte canonical id for the named built-in contribution.
+
+    ``name`` is the short name used in `builtins.rs` (e.g. ``"huffman"``,
+    ``"rans"``, ``"order1_scale_ac"``). The derivation matches
+    ``CanonicalId::derive``:
+    ``blake3(BUILTIN_PUBKEY || 0x00 || name || 0x00 || CARGO_PKG_VERSION)``
+    — the zero separators are part of the derivation so authors
+    reproducing IDs in other languages get the same bytes.
+    """
+    return bytes(_ext.builtin_canonical_id(name))
+
+
+# --- Dtype lookups (Rust-authoritative) -----------------------------------
+
+
+@translate_errors
+def dtype_element_size(dtype_code: int) -> int:
+    """Bytes per scalar element for ``dtype_code``."""
+    return _ext.dtype_element_size(dtype_code)
+
+
+# --- Container API (PPG chain-based) ------------------------------------
+
+
+@translate_errors
+def compress_model(
+    records: list[
+        tuple[
+            str,
+            int,
+            int,
+            bytes,
+            list[int] | None,
+            str | None,
+            bytes | None,
+        ]
+    ],
+    chains_per_tensor: list[list[bytes]],
+    emit_payload_hash: bool,
+    emit_plane_crc: bool,
+    method_hint: int = 3,
+    streaming_chunk: int | None = None,
+    codec_menu: list[int] | None = None,
+) -> bytes:
+    """Compress a list of tensors into a ``.ptwm`` container.
+
+    Each element of ``records`` is a 7-tuple:
+    ``(name, dtype_code, input_format, raw_bytes, shape,
+       dtype_name, delta_reference_blake3)``.
+    ``chains_per_tensor`` — one list of chain wire-byte blobs per tensor;
+    each blob is produced by ``Chain.to_bytes()`` from ``ptwm.preprocessing``.
+    Pass a single-element list when no multi-chain trial-encode is needed.
+    ``method_hint``: 1=HUFFMAN, 2=ZSTD, 3=MICROSCALE, 4=RANS, 5=IDENTITY.
+    ``streaming_chunk``: when set, split each plane into fixed-size chunks.
+    ``codec_menu``: optional list of wire-stable ``CodecId`` integers
+    restricting the trial-encode menu. ``None`` uses the full menu;
+    empty intersection on a plane raises ``InvalidContainer``.
+    """
+    return bytes(
+        _ext.compress_model(
+            records,
+            chains_per_tensor,
+            emit_payload_hash,
+            emit_plane_crc,
+            method_hint,
+            streaming_chunk,
+            codec_menu,
+        )
+    )
+
+
+@translate_errors
+def decode_tensor(blob: bytes, name: str, skip_missing: bool = False) -> bytes:
+    """Decode a single tensor from a ``.ptwm`` container blob.
+
+    Returns raw decompressed bytes.
+    """
+    return bytes(_ext.decode_tensor(blob, name, skip_missing))
+
+
+@translate_errors
+def decode_model(blob: bytes, skip_missing: bool = False) -> list[tuple[str, bytes]]:
+    """Decode every tensor from a ``.ptwm`` container blob in parallel.
+
+    Returns ``(name, raw_bytes)`` per tensor in container index order. Prefer
+    this over per-tensor :func:`decode_tensor` for a full-model load — the
+    decodes run across the rayon pool with the GIL released.
+    """
+    return _ext.decode_model(blob, skip_missing)
+
+
+@translate_errors
+def decode_tensor_shape(blob: bytes, name: str) -> tuple[list[int], str] | None:
+    """Return ``(shape, dtype_name)`` for a tensor, or ``None``."""
+    return _ext.decode_tensor_shape(blob, name)
+
+
+@translate_errors
+def decode_tensor_delta_hash(
+    blob: bytes, name: str, skip_missing: bool = False
+) -> bytes | None:
+    """Return the 32-byte BLAKE3 hash of the delta reference for ``name``.
+
+    Returns ``None`` if the tensor has no delta dependency.
+    """
+    got = _ext.decode_tensor_delta_hash(blob, name, skip_missing)
+    return bytes(got) if got is not None else None
+
+
+@translate_errors
+def list_tensor_names(blob: bytes, skip_missing: bool = False) -> list[str]:
+    """Return every tensor name stored in a ``.ptwm`` container."""
+    return list(_ext.list_tensor_names(blob, skip_missing))
+
+
+@translate_errors
+def list_plane_summaries(
+    blob: bytes,
+) -> list[tuple[str, int, int, int, int, int]]:
+    """Return per-plane summaries for every tensor in a container.
+
+    Each entry: ``(tensor_name, plane_index, role, codec_id,
+    state_source, payload_len)``.
+    """
+    return list(_ext.list_plane_summaries(blob))
+
+
+@translate_errors
+def decode_tensor_info(
+    blob: bytes, name: str, skip_missing: bool = False
+) -> tuple[list[int], str, int, int] | None:
+    """Return ``(shape, dtype_name, dtype_code, input_format)`` or ``None``."""
+    return _ext.decode_tensor_info(blob, name, skip_missing)
+
+
+# --- Container ⇄ canonical PTWM-JSON manifest transcoder -------------------
+
+
+@translate_errors
+def explode_ptwm(
+    blob: bytes,
+) -> tuple[str, list[tuple[str, str, str]], dict[str, bytes]]:
+    """Explode a ``.ptwm`` container into its canonical PTWM-JSON projection.
+
+    Returns ``(registry_json, tensors, members)`` where ``registry_json`` is
+    the canonical (JCS) archive-scoped registry manifest, ``tensors`` is a list
+    of ``(name, key, tensor_manifest_json)`` in file order, and ``members`` maps
+    member name → bytes. ``implode_ptwm`` reproduces ``blob`` byte-for-byte.
+
+    The extension already returns ``str``/``bytes`` of the correct type, so the
+    result is returned without re-wrapping (avoids copying large payloads).
+    """
+    return _ext.explode_ptwm(blob)
+
+
+@translate_errors
+def implode_ptwm(
+    registry_json: str,
+    tensor_jsons: list[str],
+    members: dict[str, bytes],
+) -> bytes:
+    """Reconstruct a ``.ptwm`` container from a manifest projection.
+
+    ``tensor_jsons`` must be in file order. Inverse of :func:`explode_ptwm`.
+    """
+    return _ext.implode_ptwm(registry_json, tensor_jsons, members)

ptwm/_rust/_errors.py

@@ -0,0 +1,38 @@
+"""Error translation at the PyO3 boundary.
+
+The `_core` extension raises stdlib `ValueError`/`RuntimeError`/`MemoryError`/
+`BufferError` subclasses (mapped from `WeightsCoreError` variants on the Rust
+side). The `translate_errors` decorator normalizes those into
+`Error` subclasses so callers can `except Error` without caring
+which Python exception type PyO3 used.
+"""
+
+from __future__ import annotations
+
+import functools
+from collections.abc import Callable
+
+from .._exceptions import Error, HeaderParseError
+
+
+def translate_errors[**P, R](fn: Callable[P, R]) -> Callable[P, R]:
+    """Wrap a callable so `_core` exceptions surface as :class:`Error`.
+
+    Header-parse errors raised by `_core.parse_header` are specifically mapped
+    to :class:`HeaderParseError`; everything else to a plain
+    :class:`Error`.
+    """
+
+    @functools.wraps(fn)
+    def wrapped(*args: P.args, **kwargs: P.kwargs) -> R:
+        try:
+            return fn(*args, **kwargs)
+        except Error:
+            raise
+        except (BufferError, MemoryError, ValueError, RuntimeError) as e:
+            message = str(e)
+            if fn.__name__ == "parse_header":
+                raise HeaderParseError(message) from e
+            raise Error(message) from e
+
+    return wrapped

ptwm/_rust/ext.py

@@ -0,0 +1,19 @@
+"""Thin shim exposing the install/discovery surface as ptwm._rust.ext."""
+
+from __future__ import annotations
+
+from ptwm._core import (  # type: ignore[import]
+    InstalledContribution,
+    ext_install,
+    ext_list,
+    ext_remove,
+    ext_user_dir,
+)
+
+__all__ = [
+    "InstalledContribution",
+    "ext_install",
+    "ext_list",
+    "ext_remove",
+    "ext_user_dir",
+]

ptwm/_rust/host.py

@@ -0,0 +1,7 @@
+"""Thin shim exposing the host-flavor registry as ptwm._rust.host."""
+
+from __future__ import annotations
+
+from ptwm._core import ContributionRegistry, host_registry_entries
+
+__all__ = ["ContributionRegistry", "host_registry_entries"]

ptwm/_rust/policy.py

@@ -0,0 +1,7 @@
+"""Thin shim re-exporting policy types from the Rust extension."""
+
+from __future__ import annotations
+
+from ptwm._core import PolicyFile, ResolvedPolicy  # type: ignore[import]
+
+__all__ = ["PolicyFile", "ResolvedPolicy"]