Skip to content

Refactor trap workflow API: add randomize_model and propagate permuted_ids#370

Open
charlesmartin14 wants to merge 7 commits intomasterfrom
codex/rework-api-for-trap-analysis-and-metrics
Open

Refactor trap workflow API: add randomize_model and propagate permuted_ids#370
charlesmartin14 wants to merge 7 commits intomasterfrom
codex/rework-api-for-trap-analysis-and-metrics

Conversation

@charlesmartin14
Copy link
Copy Markdown
Member

Motivation

  • Decouple the trap workflow so permutation/randomization can be performed independently and reused to avoid re-randomization, sync issues, and repeated work across analysis/removal stages.
  • Allow downstream trap analysis and removal to accept a pre-randomized model and deterministic permuted_ids to guarantee consistent trap indexing and fingerprints across tools.

Description

  • Added randomize_model(ww, model, ...) in weightwatcher/remove_traps.py which deep-copies and permutes eligible layer weights and returns (randomized_model, permuted_ids).
  • Exposed a public WeightWatcher.randomize_model(...) method that forwards to the new helper for callers to obtain a randomized model and per-layer permutation ids.
  • Extended WeightWatcher.analyze_traps(...) signature to accept randomized_model and permuted_ids, and updated trap_analysis.analyze_traps(...) to accept and forward permuted_ids into per-layer params.
  • Updated apply_analyze_traps to use externally supplied permuted_ids when present instead of generating new permutations so analysis and removal remain synchronized.
  • Files changed: weightwatcher/remove_traps.py, weightwatcher/weightwatcher.py, and weightwatcher/trap_analysis.py.

Testing

  • Ran targeted unit tests: pytest -q tests/test_analyze_traps.py::TestAnalyzeTraps::test_analyze_traps_trap_burden_columns_appear tests/test_remove_traps.py::test_remove_traps_public_api_direct_call which produced 1 passed, 1 skipped.
  • No other automated tests were run in this change set.

Codex Task

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant