"""memory-native — finite-state counter synapses + reversible activations, in pure PyTorch.

A training method that attacks all four memory pools of training (params, optimizer state,
gradients, activations) with two independent levers, implemented on stock PyTorch (CPU/CUDA),
with no custom engine:

  * CompactCounterLinear / RMSCounterLinear -- a ternary weight whose optimizer lives inside
    a per-synapse finite-state automaton; the update is fused into backward (no FP master, no
    Adam moments, no full gradient buffer).
  * ReversibleCouplingBlock -- activations are recomputed in backward instead of stored.

See README for what is measured vs what still needs the Triton/CUDA packed kernel.
"""
from .counter import (
    C_DEFAULT,
    CompactCounterLinear,
    RMSCounterLinear,
    decode_state,
    encode_state,
    stochastic_round,
    ternary_gradient_unbiased,
)
from .reversible import ReversibleCouplingBlock, ReversibleSequence, ReversibleSequential
from .packed import PackedRMSCounterLinear, pack_codes, unpack_codes
from .slowfast import SlowFastCounterLinear
from .triton_counter import HAS_TRITON, TritonCounterLinear, triton_decode_matmul, triton_grad_x
from .baselines import TernaryQATLinear, make_linear
from .models import CONFIGS, GPT, GPTConfig
from .reversible_gpt import ReversibleGPT
from .fused_qkv import CounterQKVLinear
from .memory_ffn import CounterMemoryFFN, CounterValueMemory
from .moe_ffn import CounterMoEFFN
from .mtp import MultiTokenHead, MTPGPT, mtp_loss, shift_targets
from .group_counter import GroupCounterLinear
from .stack_linear import StackCounterLinear
from .mod import MoDBlock
from .glm import MNGLM, ReversibleMNGLM, GLMBlock, GLMAttention, RMSNorm
from .int8_compute import (
    fp8_correlation,
    int4_correlation,
    int8_correlation,
    int8_correlation_presaved,
    int8_forward_ternary,
    int8_mm,
    quantize_int4_cols,
    quantize_int8_cols,
    quantize_int8_rows,
)
from .memory import compare_training_peak, fmt_bytes, memory_report, peak_training_memory
from .actquant import effective_bits, quantize_codes, stochastic_quantize
from .budget import BudgetRow, format_budget, training_budget
from .optimizers import GaLoreAdamW, LoMo, available_optimizers, build_optimizer
from .data import get_batch, load_corpus, synthetic_corpus

__version__ = "0.1.0"

__all__ = [
    "C_DEFAULT",
    "CompactCounterLinear",
    "RMSCounterLinear",
    "SlowFastCounterLinear",
    "PackedRMSCounterLinear",
    "pack_codes",
    "unpack_codes",
    "TritonCounterLinear",
    "triton_decode_matmul",
    "triton_grad_x",
    "HAS_TRITON",
    "encode_state",
    "decode_state",
    "stochastic_round",
    "ternary_gradient_unbiased",
    "ReversibleCouplingBlock",
    "ReversibleSequential",
    "ReversibleSequence",
    "TernaryQATLinear",
    "make_linear",
    "GPT",
    "ReversibleGPT",
    "CounterQKVLinear",
    "CounterMemoryFFN",
    "CounterValueMemory",
    "CounterMoEFFN",
    "MultiTokenHead",
    "MTPGPT",
    "mtp_loss",
    "shift_targets",
    "GroupCounterLinear",
    "StackCounterLinear",
    "MoDBlock",
    "MNGLM",
    "ReversibleMNGLM",
    "GLMBlock",
    "GLMAttention",
    "RMSNorm",
    "fp8_correlation",
    "int4_correlation",
    "int8_correlation",
    "int8_correlation_presaved",
    "int8_forward_ternary",
    "int8_mm",
    "quantize_int4_cols",
    "quantize_int8_cols",
    "quantize_int8_rows",
    "GPTConfig",
    "CONFIGS",
    "memory_report",
    "peak_training_memory",
    "compare_training_peak",
    "fmt_bytes",
    "build_optimizer",
    "available_optimizers",
    "GaLoreAdamW",
    "LoMo",
    "stochastic_quantize",
    "quantize_codes",
    "effective_bits",
    "training_budget",
    "format_budget",
    "BudgetRow",
    "load_corpus",
    "get_batch",
    "synthetic_corpus",
    "__version__",
]