Comparator fix for weakref

[aseembits93]

Some extra tests to ensure more torch datatypes are supported but no changes in the comparator

[claude]

PR Review Summary

Prek Checks

✅ All checks pass — ruff check and ruff format both passed with no issues.

Mypy Analysis

⚠️ Pre-existing mypy errors (27 in comparator.py, ~399 in test_comparator.py) — all are pre-existing issues (mostly no-any-return and missing return type annotations on test methods). No new mypy errors introduced by this PR.

The PR actually improves typing by adding superset_obj: bool = False type annotation to the comparator() function signature.

Code Review

✅ No critical issues found. This is a clean, well-structured PR.

Changes reviewed:

• codeflash/verification/comparator.py: Added weakref handling in the comparator (lines 175-184). The implementation correctly:
  • Checks isinstance(orig, weakref.ref) after the type-equality check at line 127, so new is guaranteed to also be a weakref.ref
  • Handles dead references (both dead = equal, one dead = not equal)
  • Recursively compares referents
• tests/test_comparator.py: Added comprehensive test coverage for:
  • weakref.ref objects (same object, equivalent objects, different objects, dead refs, nested structures, in containers)
  • Custom objects with weakrefs and callbacks
  • torch.nn modules: Linear, Conv2d, BatchNorm, Dropout, activations, pooling, Embedding, LSTM, GRU, LayerNorm, MultiheadAttention, Sequential, ModuleList, ModuleDict, custom modules, nested modules, Transformer, parameter/buffer modifications, device placement, Conv1d/Conv3d, Flatten/Unflatten, Identity, superset mode

Test Coverage

File	Stmts	Miss	Coverage
codeflash/verification/comparator.py	341	130	62%
tests/test_comparator.py	3209	1269	60%

New code coverage: The new weakref handling code (lines 175-184) is fully covered by the new tests. All 3 weakref test functions pass.

Missing lines are all pre-existing uncovered paths for optional dependencies (JAX, xarray, TensorFlow, scipy, torch, numba, pyrsistent, sqlalchemy) that are not available in the test environment. Many torch.nn tests are skipped (42 skipped) because torch is not installed in the test runner.

Coverage assessment: No coverage regression. New code is fully tested. The 62% file coverage is expected given the many optional library branches.

Codeflash Optimization PRs

No optimization PRs are eligible for merge — all open codeflash-ai[bot] PRs have failing CI checks.

---

Last updated: 2026-02-06