add utility to compute analytical bits per weight

[pkmandke]

Add utility to compute the average bits-per-weight (bpw) of a prepared coreai-opt model in eager mode quantization and palettization.

Public API

Imported from coreai_opt.inspection:

def bits_per_weight(model: torch.nn.Module) -> BitsPerWeightResult

@dataclass
class BitsPerWeightResult:
    bpw: float                      # total_bits / total_weights, 0.0 if no params
    per_module: dict[str, float]    # module name -> module's own average bpw
    total_bits: int                 # total storage cost in bits, incl. amortized overhead
    total_weights: int              # total logical parameter elements

Usage:

from coreai_opt.inspection import BitsPerWeightResult, bits_per_weight

result = bits_per_weight(prepared_model)
result.bpw          # 4.13
result.per_module   # {"l1": 4.06, "l2": 4.25, ...}

Supported: eager-mode integer weight quantization (int8/int4/int2 and unsigned variants, symmetric or asymmetric, any granularity, sub-byte payloads packed at n_bits), and palettization at any spec-supported n_bits including a quantized LUT.

Raises NotImplementedError for: graph-mode / torch.fx.GraphModule models, floating-point (FP8/FP4) weight quantization, and parametrizations that store multiple original tensors (weight_norm, spectral_norm).

Note: This is an analytical estimate, not a measurement. It is meant for prepared models, not finalized ones.

Testing

• tests/inspection/test_bits_per_weight.py: unit tests against hand-derived golden bit counts across quantization dtypes, qschemes, and granularities, plus palettization n_bits x granularity combinations. Also covers persistent and non-persistent buffers, tied weights, per-module attribution, and the unsupported-config errors.
• tests/export/test_bpw_export_size.py: cross-checks total_bits / 8 against the measured payload of an actual Core AI export. The prediction must be a lower bound and land within a 2.5 percent structural-metadata budget.
• tests/export/export_utils.py: adds coreai_export_size_bytes helper.
• tests/models/simple.py: adds LinearBatchNormModel fixture (buffer coverage) and optional bias to two existing fixtures.

---

TODO:

• Add export tests to validate analytical bpw against exported asset size
• Add tests for per-module bpw

[u-simha]

Have left some comments; my main suggestion would be to have a documentation page (even if it is brief) rather than the doc string at the top of the file

[u-simha]

My suggestion would be to have a doc page explaining the usage of this tool; doc strings in the headings of the file isn't often read, compared to the documentation.