src/agiws_neural_quant/__init__.py

"""AGIWSNeuralQuant — Universal neural network quantization library.

Unified architecture: one parameterized Quantizer for all formats,
QuantizedWeight/QuantizedActivation containers, dual-path cross-quantization
distillation, chunked dequant for minimal VRAM, QAT with learnable parameters
via STE, SSM-aware exclusion patterns.

Quantization primitives (pure tensor-level functions, no nn.Module wrappers):
  - ternary (BitNet 1.58): ternarize_tensor, ternary_dequantize, fake_ternarize
  - nf (QLoRA NormalFloat generalized): NF2/NF3/NF4/NF8 LUTs, quantize/dequantize, pack/unpack, double_quant
  - fp8 (E4M3/E5M2): FP8 LUTs, quantize/dequantize
  - fp4 (E2M1 / NVFP4 / MXFP4): FP4 LUT, quantize/dequantize, pack/unpack, E8M0
  - ste (Straight-Through Estimator): STEQuantize, fake_quantize

Layer wrappers live ONLY in base.py (QuantizedModule) — no per-format
nn.Module classes. The unified Quantizer (quantizer.py) + presets (presets.py)
configure all 40+ formats as parameters, not separate classes.
"""

__version__ = "0.6.5"

# Core unified architecture
from agiws_neural_quant.dispatch import quantize_model, count_quantizable_layers, make_quantizer, save_model, load_model
from agiws_neural_quant.base import QuantizedModule, QuantizedWeight, QuantizedActivation
from agiws_neural_quant.quantizer import Quantizer
from agiws_neural_quant.presets import FORMAT_PRESETS, get_preset

# SSM-aware exclusion patterns (KDA / Mamba / RWKV / linear attention)
from agiws_neural_quant.ssm_patterns import (
    get_ssm_exclude_patterns,
    get_ssm_subtree_patterns,
)

# QAT + dual-path distillation (unified)
from agiws_neural_quant.training_unified import (
    UnifiedQATWrapper,
    dual_path_loss,
    strip_latent,
)

# Quantization primitives (pure functions / LUTs — no nn.Module wrappers)
from agiws_neural_quant.training import fake_quantize, STEQuantize
from agiws_neural_quant.training.ste import STECodebook, fake_codebook_quantize
from agiws_neural_quant.ternary import (
    ternarize_tensor,
    ternary_dequantize,
    fake_ternarize,
)
from agiws_neural_quant.nf import (
    make_normalfloat_lut,
    NF2_LUT,
    NF3_LUT,
    NF4_LUT,
    NF8_LUT,
    quantize_nf,
    dequantize_nf,
    pack_nf,
    unpack_nf,
    double_quantize_scales_2d,
    dequantize_scales_2d,
)
from agiws_neural_quant.fp8 import (
    FP8_E4M3_LUT,
    FP8_E5M2_LUT,
    quantize_fp8,
    dequantize_fp8,
)
from agiws_neural_quant.fp4 import (
    FP4_E2M1_LUT,
    quantize_fp4,
    dequantize_fp4,
    pack_fp4,
    unpack_fp4,
    E8M0_LUT,
)
from agiws_neural_quant.fp6 import (
    FP6_E3M2_LUT,
    FP6_E2M3_LUT,
    quantize_fp6,
    dequantize_fp6,
    pack_fp6,
    unpack_fp6,
)
from agiws_neural_quant import kquant
from agiws_neural_quant.kquant import (
    quantize_blocks,
    dequantize_blocks,
)

# Subsystem: teacher-cache for distillation
from agiws_neural_quant.cache import (
    get_cache_dir,
    get_cache_path,
    is_cache_valid,
    get_sources_needing_cache,
    save_cache,
    load_cache,
    CaptureConfig,
    select_modules,
    capture_with_hooks,
    save_layer_cache,
    load_layer_cache,
    load_layer_io,
    list_cached_names,
    validate_cache_contents,
    clean_old_cache,
    TeacherCache,
)

# Subsystem: checkpoint extraction (vision encoders, submodules, shards)
from agiws_neural_quant.extract import (
    extract_subcheckpoint,
    extract_vision_encoder,
    extract_module_group,
    load_subcheckpoint,
    list_shards,
    find_keys,
    ExtractReport,
)

# Subsystem: universal file-to-file converter (low memory)
from agiws_neural_quant.converters import (
    convert_model,
    detect_format,
    list_safetensors_keys,
    read_safetensors_tensor,
    stream_safetensors,
    write_safetensors,
    detect_quant_layout,
    GGUFReader,
    load_gguf_modules,
    gguf_to_quantized_modules,
    convert_nvfp4_safetensors_tensor,
    dequantize_nvfp4_safetensors,
)

# GGUF k-quant packed slice-dequant (on-the-fly chunked, llama.cpp-style VRAM)
from agiws_neural_quant.kquant.gguf_packed import (
    dequant_gguf_slice,
    bytes_per_row as gguf_bytes_per_row,
)

# Transformers integration — NeuralQuant as quantization plugin for from_pretrained
from agiws_neural_quant.transformers_integration import (
    NeuralQuantConfig,
    NeuralQuantHfQuantizer,
    NEURAL_QUANT_METHOD,
)

# bitsandbytes packed format reader (read bnb safetensors without bnb library)
from agiws_neural_quant.bnb.bnb_packed import (
    dequant_bnb_slice,
    dequant_bnb_nf4,
    dequant_bnb_int8,
)
from agiws_neural_quant.bnb.bnb_loader import load_bnb_model

# Subsystem: layer analysis (per-layer quantization suitability)
from agiws_neural_quant.analysis import LayerAnalyzer, SplitReport, LayerResult, AnomalyReport

# Subsystem: TradingLR scheduler (auto-LR on technical indicators — separate concern)
from agiws_neural_quant.trading_lr import TradingLR, DEMA, ATR, RSI

# Subsystem: architectural replacement via dual-path QAT (Jamba-conversion, Stage 19e)
from agiws_neural_quant.arch_replace import MambaLayer, ReplaceModule, jamba_replace

# Subsystem: Multi-Token Prediction head via dual-path QAT (Stage 24)
from agiws_neural_quant.mtp import MTPHead, attach_mtp

__all__ = [
    # Core
    "quantize_model",
    "count_quantizable_layers",
    "make_quantizer",
    "save_model",
    "load_model",
    "QuantizedModule",
    "QuantizedWeight",
    "QuantizedActivation",
    "Quantizer",
    "FORMAT_PRESETS",
    "get_preset",
    # SSM patterns
    "get_ssm_exclude_patterns",
    "get_ssm_subtree_patterns",
    # QAT + dual-path
    "UnifiedQATWrapper",
    "dual_path_loss",
    "strip_latent",
    # Primitives: STE
    "fake_quantize",
    "STEQuantize",
    "STECodebook",
    "fake_codebook_quantize",
    # Primitives: ternary (BitNet 1.58)
    "ternarize_tensor",
    "ternary_dequantize",
    "fake_ternarize",
    # Primitives: NormalFloat (NF2/NF3/NF4/NF8 — QLoRA generalized)
    "make_normalfloat_lut",
    "NF2_LUT",
    "NF3_LUT",
    "NF4_LUT",
    "NF8_LUT",
    "quantize_nf",
    "dequantize_nf",
    "pack_nf",
    "unpack_nf",
    "double_quantize_scales_2d",
    "dequantize_scales_2d",
    # Primitives: FP8
    "FP8_E4M3_LUT",
    "FP8_E5M2_LUT",
    "quantize_fp8",
    "dequantize_fp8",
    # Primitives: FP4 / NVFP4 / MXFP4
    "FP4_E2M1_LUT",
    "quantize_fp4",
    "dequantize_fp4",
    "pack_fp4",
    "unpack_fp4",
    "E8M0_LUT",
    # Primitives: FP6 (E3M2 / E2M3)
    "FP6_E3M2_LUT",
    "FP6_E2M3_LUT",
    "quantize_fp6",
    "dequantize_fp6",
    "pack_fp6",
    "unpack_fp6",
    # Primitives: kquant (GGUF k-quants super-block layout)
    "kquant",
    "quantize_blocks",
    "dequantize_blocks",
    # Teacher cache
    "get_cache_dir",
    "get_cache_path",
    "is_cache_valid",
    "get_sources_needing_cache",
    "save_cache",
    "load_cache",
    "CaptureConfig",
    "select_modules",
    "capture_with_hooks",
    "save_layer_cache",
    "load_layer_cache",
    "load_layer_io",
    "list_cached_names",
    "validate_cache_contents",
    "clean_old_cache",
    "TeacherCache",
    # Extraction
    "extract_subcheckpoint",
    "extract_vision_encoder",
    "extract_module_group",
    "load_subcheckpoint",
    "list_shards",
    "find_keys",
    "ExtractReport",
    # Universal converter
    "convert_model",
    "detect_format",
    "list_safetensors_keys",
    "read_safetensors_tensor",
    "stream_safetensors",
    "write_safetensors",
    "detect_quant_layout",
    "GGUFReader",
    "load_gguf_modules",
    "gguf_to_quantized_modules",
    "convert_nvfp4_safetensors_tensor",
    "dequantize_nvfp4_safetensors",
    # GGUF k-quant packed slice-dequant
    "dequant_gguf_slice",
    "gguf_bytes_per_row",
    # Transformers integration (from_pretrained quantization plugin)
    "NeuralQuantConfig",
    "NeuralQuantHfQuantizer",
    # bitsandbytes reader
    "dequant_bnb_slice",
    "dequant_bnb_nf4",
    "dequant_bnb_int8",
    "load_bnb_model",
    # Layer analysis
    "LayerAnalyzer",
    "SplitReport",
    "LayerResult",
    "AnomalyReport",
    # TradingLR (separate subsystem)
    "TradingLR",
    "DEMA",
    "ATR",
    "RSI",
    # Architectural replacement (Stage 19e — Jamba via dual-path QAT)
    "MambaLayer",
    "ReplaceModule",
    "jamba_replace",
    # MTP head (Stage 24 — Multi-Token Prediction via dual-path QAT)
    "MTPHead",
    "attach_mtp",
]