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
  - nf4 (QLoRA): NF4_LUT, quantize_nf4, dequantize_nf4, 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.3.0"

# 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.nf4 import (
    NF4_LUT,
    quantize_nf4,
    dequantize_nf4,
    pack_nf4,
    unpack_nf4,
    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,
    convert_nvfp4_safetensors_tensor,
    dequantize_nvfp4_safetensors,
)

# 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

__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: NF4 (QLoRA)
    "NF4_LUT",
    "quantize_nf4",
    "dequantize_nf4",
    "pack_nf4",
    "unpack_nf4",
    "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",
    "convert_nvfp4_safetensors_tensor",
    "dequantize_nvfp4_safetensors",
    # Layer analysis
    "LayerAnalyzer",
    "SplitReport",
    "LayerResult",
    "AnomalyReport",
    # TradingLR (separate subsystem)
    "TradingLR",
    "DEMA",
    "ATR",
    "RSI",
]