sweep

MIP sweep functionality for exploring multiple memory compression rates.

Functions

`get_teacher_memory_from_subblock_stats`	Calculate teacher model memory from subblock_stats.json.
`get_teacher_num_params_from_subblock_stats`	Calculate total teacher parameter count from subblock_stats.json.
`extract_solution_results`	Extract results from a completed MIP solution.
`write_results_to_csv`	Write sweep results to CSV file.
`run_mip_sweep`	Run MIP for multiple memory compression rates and generate CSV with results.

extract_solution_results(solution_path, target_memory_mib, teacher_memory_mib, compression_rate)

Extract results from a completed MIP solution.

Parameters:

solution_path (Path) – Path to the solutions.json file (not the directory!)
target_memory_mib (float) – Target memory constraint used for MIP
teacher_memory_mib (float) – Teacher model memory in MiB
compression_rate (float) – Compression rate applied

Returns:

Dictionary containing extracted metrics

Return type:

dict

get_teacher_memory_from_subblock_stats(hydra_cfg)

Calculate teacher model memory from subblock_stats.json.

Sums non_block and per-layer _get_block_stats(subblock_stats, block_config, layer_index) over model_config.block_configs, matching run_puzzle._get_block_stats().

Parameters:: hydra_cfg (DictConfig) – Hydra configuration object
Returns:: Total teacher memory in MiB
Return type:: float

get_teacher_num_params_from_subblock_stats(hydra_cfg)

Calculate total teacher parameter count from subblock_stats.json.

Sums non_block and per-layer _get_block_stats(...)["num_params"] over model_config.block_configs, matching run_puzzle._get_block_stats().

Parameters:: hydra_cfg (DictConfig) – Hydra configuration object
Returns:: Total teacher parameter count (same units as subblock_stats JSON).
Return type:: int

run_mip_sweep(hydra_cfg)

Run MIP for multiple memory compression rates and generate CSV with results.

This function is called when mip.sweep.enabled is True in the config.

Parameters:: hydra_cfg – Hydra configuration object with mip.sweep settings

write_results_to_csv(results, output_csv)

Write sweep results to CSV file.

Parameters:

results (list) – List of result dictionaries
output_csv (str) – Path to output CSV file