module ptwm.preprocessing

def cache_dir() -> Path

def cache_info() -> dict[str, int]

def cache_path_for(dtype_code: int, role: ClassifierRole) -> Path

class CacheEntry:

def __init__(self, internal_dtype: int, ops: tuple[int, ...], byte_split_planes: int) -> None

byte_split_planes: int

internal_dtype: int

ops: tuple[int, ...]

def to_builder(self) -> ChainBuilder

class Chain:

CHAIN_BF16_SPLIT: ChainBuilder = _bf16_split_chain

CHAIN_BYTE_PASSTHROUGH_VALUE: ChainBuilder = _u8_byte_passthrough_value_chain

CHAIN_F32_GLOBAL_SCALE: ChainBuilder = _fp32_global_scale_chain

CHAIN_F8E4M3_NIBBLE: ChainBuilder = _fp8e4m3_nibble_chain

CHAIN_F8E4M3_SCALE_PASSTHROUGH: ChainBuilder = _fp8e4m3_scale_passthrough_chain

CHAIN_FP32_SPLIT: ChainBuilder = _fp32_split_chain

CHAIN_FP4_NIBBLE_PASSTHROUGH: ChainBuilder = _fp4_passthrough_chain

CHAIN_FP8E4M3_NIBBLE: ChainBuilder = _fp8e4m3_nibble_chain

CHAIN_FP8E5M2_NIBBLE: ChainBuilder = _fp8e5m2_nibble_chain

CHAIN_U8_SCALE_PASSTHROUGH: ChainBuilder = _u8_scale_passthrough_chain

def __init__(self, nodes: list[ChainNode], edges: list[ChainEdge], terminals: list[TerminalRef]) -> None

edges: list[ChainEdge]

nodes: list[ChainNode]

terminals: list[TerminalRef]

def to_bytes(self) -> bytes

ChainBuilder = Callable[[list[int]], Chain | None]

class ChainEdge:

def __init__(self, src_node: int, src_output_idx: int, dst_node: int, dst_input_idx: int, role_override: bytes | None = None, vendor_bytes: bytes = b'') -> None

dst_input_idx: int

dst_node: int

role_override: bytes | None = None

src_node: int

src_output_idx: int

vendor_bytes: bytes = field(default=b'')

class ChainNode:

def __init__(self, op_id: int, params: bytes = b'') -> None

op_id: int

params: bytes = field(default=b'')

def chains_for(dtype_code: int, role: ClassifierRole) -> list[ChainBuilder]

class ClassifierRole(IntEnum):

PACKED_VALUES = 3

SCALE_BLOCK = 1

SCALE_GLOBAL = 2

STANDARD = 0

def clear_cache() -> int

def explore_chains(dtype_code: int, role: ClassifierRole, opts: ExploreOptions | None = None) -> Iterator[ChainBuilder]

class ExploreOptions:

def __init__(self, max_nodes: int = 8, max_depth: int = 5, max_candidates_per_dtype: int = 32, allow_cross_tensor: bool = False, allow_cross_plane: bool = False, time_budget_ms: int = 5000) -> None

allow_cross_plane: bool = False

allow_cross_tensor: bool = False

max_candidates_per_dtype: int = 32

max_depth: int = 5

max_nodes: int = 8

time_budget_ms: int = 5000

def is_legal_chain(chain: Chain, opts: ExploreOptions | None = None) -> bool

def load_cached_builders(dtype_code: int, role: ClassifierRole) -> list[ChainBuilder]

def load_cached_entries(dtype_code: int, role: ClassifierRole) -> list[CacheEntry]

class Plane:

def __init__(self, id: int, role: str, data: bytes) -> None

data: bytes

id: int

role: str

PRODUCTION_CHAINS: dict[tuple[int, ClassifierRole], list[ChainBuilder]] = {(_DT.FLOAT8_E4M3FN.code, ClassifierRole.SCALE_BLOCK): [CHAIN_F8E4M3_SCALE_PASSTHROUGH, CHAIN_F8E4M3_NIBBLE], (_DT.UINT8.code, ClassifierRole.SCALE_BLOCK): [CHAIN_U8_SCALE_PASSTHROUGH], (_DT.FLOAT32.code, ClassifierRole.SCALE_GLOBAL): [CHAIN_F32_GLOBAL_SCALE], (_DT.FLOAT4_E2M1FN_X2.code, ClassifierRole.PACKED_VALUES): [CHAIN_FP4_NIBBLE_PASSTHROUGH], (_DT.BFLOAT16.code, ClassifierRole.STANDARD): [CHAIN_BF16_SPLIT, CHAIN_BF16_SPLIT_INTDELTA, CHAIN_BF16_SPLIT_PREDICTOR_XOR, CHAIN_BF16_MICROSCALE_REPACK, CHAIN_BF16_SPLIT_MZS_K1, CHAIN_BF16_SPLIT_MZS_K2, CHAIN_BF16_SPLIT_MZS_K3, CHAIN_BF16_SPLIT_MZS_K4, CHAIN_BF16_SPHERICAL_NORMALIZE, CHAIN_BF16_ALPHA_STABLE_NORMALIZE, CHAIN_BF16_HUFF_LLM_RAW], (_DT.FLOAT16.code, ClassifierRole.STANDARD): [CHAIN_BF16_SPLIT, CHAIN_BF16_SPLIT_INTDELTA, CHAIN_BF16_SPLIT_PREDICTOR_XOR, CHAIN_BF16_SPLIT_MZS_K1, CHAIN_BF16_SPLIT_MZS_K2, CHAIN_BF16_SPLIT_MZS_K3, CHAIN_BF16_SPLIT_MZS_K4, CHAIN_FP16_SPHERICAL_NORMALIZE, CHAIN_FP16_ALPHA_STABLE_NORMALIZE, CHAIN_FP16_HUFF_LLM_RAW], (_DT.FLOAT32.code, ClassifierRole.STANDARD): [CHAIN_FP32_SPLIT, CHAIN_FP32_SPLIT_INTDELTA, CHAIN_FP32_SPLIT_PREDICTOR_XOR, CHAIN_FP32_SPLIT_MZS_K1, CHAIN_FP32_SPLIT_MZS_K2, CHAIN_FP32_SPLIT_MZS_K3, CHAIN_FP32_SPLIT_MZS_K4, CHAIN_FP32_SPHERICAL_NORMALIZE, CHAIN_FP32_ALPHA_STABLE_NORMALIZE], (_DT.FLOAT8_E4M3FN.code, ClassifierRole.STANDARD): [CHAIN_FP8E4M3_NIBBLE], (_DT.FLOAT8_E5M2.code, ClassifierRole.STANDARD): [CHAIN_FP8E5M2_NIBBLE], (_DT.UINT8.code, ClassifierRole.STANDARD): [CHAIN_BYTE_PASSTHROUGH_VALUE]}

def save_cached_entries(dtype_code: int, role: ClassifierRole, new_entries: Iterable[CacheEntry]) -> int

class TerminalRef:

def __init__(self, node_idx: int, output_idx: int, role: bytes) -> None

node_idx: int

output_idx: int

role: bytes