Skip to content

refactor(ata): move target-gain arithmetic to Rust - #805

Merged
seonghobae merged 13 commits into
mainfrom
fix/ata-target-gain-rust
Aug 12, 2026
Merged

refactor(ata): move target-gain arithmetic to Rust#805
seonghobae merged 13 commits into
mainfrom
fix/ata-target-gain-rust

Conversation

@seonghobae

@seonghobae seonghobae commented Aug 12, 2026

Copy link
Copy Markdown
Contributor

Why

Protected main still computes the result-affecting capped-shortfall gain for target-information ATA in Python/NumPy. PR #804 attempted to accelerate that Python numerical owner with broadcasting, but the repository contract and issue #609 require production item-bank selection-objective arithmetic to be Rust-owned.

Intentional RED

Exact fail-first head c6acf7279f8b2d56c98857d78c753defdbbe5e97 adds only a public-boundary regression. It injects a fake compiled ATA core and requires assemble_to_target() to delegate target-gain arithmetic through that boundary. Current protected-main behavior completes assembly but never calls the injected core, so the expected RED is the ownership assertion rather than setup/import/fixture failure.

Minimum GREEN

  • add a bounded Rust target-information gain kernel over the precomputed item-information matrix;
  • expose it through a dedicated PyO3 secondary module following the existing rotation/bifactor modular-entrypoint pattern;
  • keep Python responsible for validation, candidate/content/exposure orchestration and deterministic tie breaking while Rust owns the capped-shortfall numerical objective;
  • avoid candidate-by-point broadcast temporaries and per-call thread creation; this kernel is intentionally serial until a benchmark proves parallel/GPU execution is material;
  • require direct Rust tests plus public Python parity/ownership evidence and existing ATA behavioral tests;
  • add changelog/doctoring after GREEN; no formula, model family, dependency, hosted state, workflow or release-version change.

Advances the Rust-first ATA/item-bank numerical boundary in #609 and preserves the performance concern from closed #804 without deepening the Python numerical kernel.

Summary by CodeRabbit

  • Performance

    • Improved ATA target-information gain calculations for faster, more memory-efficient candidate evaluation.
  • Reliability

    • Added validation for invalid dimensions, values, targets, accumulated information, and candidate indices.
    • Invalid inputs now produce clear value errors.
  • Compatibility

    • Improved coexistence of ATA and other native extension modules.
  • Tests

    • Added coverage for calculation accuracy, satisfied targets, invalid inputs, delegation, and extension loading.

@coderabbitai

coderabbitai Bot commented Aug 12, 2026

Copy link
Copy Markdown
Contributor

Review Change Stack

Important

Review skipped

Auto incremental reviews are disabled on this repository.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro Plus

Run ID: c06e1f32-62af-4e66-ab13-ea4b3ff0c93d

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

  • 🔍 Trigger review
📝 Walkthrough

Walkthrough

The change adds a Rust/PyO3 ATA gain kernel, validates inputs, loads the secondary extension from Python, and delegates candidate scoring from ATA assembly. Tests cover arithmetic parity, delegation, invalid inputs, and extension registration.

Changes

ATA Rust gain integration

Layer / File(s) Summary
Validated Rust gain kernel
crates/fast-mlsirm-py/src/ata_bindings.rs, crates/fast-mlsirm-py/src/entrypoint.rs
The _ata_core module validates ATA inputs, computes capped-shortfall gains, exposes target_information_gains, and registers the module. Rust tests cover valid and invalid inputs.
Python loader and assembly wiring
python/fast_mlsirm/_ata_core_loader.py, python/fast_mlsirm/ata.py
Python loads and caches the shared _ata_core extension. ATA assembly delegates candidate gain calculations to the Rust function.
Delegation and extension regression coverage
tests/test_ata.py, tests/test_modular_extension_coexistence.py, docs/changelog.d/805-ata-target-gain-rust.md
Tests verify delegation, reference-formula parity, invalid candidate handling, and registration of secondary extension modules. The changelog documents the implementation split.

Estimated code review effort: 3 (Moderate) | ~20 minutes

Sequence Diagram(s)

sequenceDiagram
  participant ATAAssembly
  participant ata_core_loader
  participant _core
  participant _ata_core
  participant RustGainKernel
  ATAAssembly->>ata_core_loader: ata_core()
  ata_core_loader->>_core: locate shared library
  ata_core_loader->>_ata_core: initialize extension
  ATAAssembly->>_ata_core: target_information_gains(matrix, cand, target, accum)
  _ata_core->>RustGainKernel: validate inputs and calculate gains
  RustGainKernel-->>_ata_core: capped-shortfall gains
  _ata_core-->>ATAAssembly: gain array
Loading

Possibly related PRs

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name Status Explanation Resolution
Docstring Coverage ⚠️ Warning Docstring coverage is 78.57% which is insufficient. The required threshold is 80.00%. Write docstrings for the functions missing them to satisfy the coverage threshold.
✅ Passed checks (4 passed)
Check name Status Explanation
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Title check ✅ Passed The title clearly and concisely describes moving ATA target-gain arithmetic to Rust.
Linked Issues check ✅ Passed The PR addresses issue #804 by moving ATA target-gain computation to Rust while preserving numerical results and scoring behavior.
Out of Scope Changes check ✅ Passed The changes support the Rust ATA kernel, integration, tests, and documentation without unrelated scope.
✨ Finishing Touches 💡 1
📝 Generate docstrings 💡
  • Create stacked PR
  • Commit on current branch
🧪 Generate unit tests (beta)
  • Create PR with unit tests
  • Commit unit tests in branch fix/ata-target-gain-rust

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

Comment @coderabbitai help to get the list of available commands.

@seonghobae
seonghobae marked this pull request as ready for review August 12, 2026 10:18

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 3

🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@crates/fast-mlsirm-py/src/ata_bindings.rs`:
- Around line 76-78: In the ATA binding flow around information_matrix.shape()
and candidate_indices, validate the matrix dimensions and documented ATA limits
before converting candidates or allocating any output buffer. Enforce a maximum
candidate count on the public input before candidate_values.as_slice() and
Vec::with_capacity, while preserving the existing empty-matrix rejection
behavior.
- Around line 71-74: Update py_target_information_gains to normalize PyO3/NumPy
input extraction failures into PyValueError, including wrong-dtype, non-array,
and non-contiguous arguments, instead of allowing TypeError to escape. Validate
the candidate count against the supported limit before Vec allocation or
candidate conversion. Add coverage for wrong-dtype, non-array, non-contiguous,
and oversized-candidate inputs.

In `@python/fast_mlsirm/ata.py`:
- Line 344: Replace the unconditional ata_core() call in the ATA assembly flow
with the established {numpy,rust,auto} backend resolver, retaining Rust as the
default when available and selecting the numerically identical NumPy reference
implementation when the Rust extension is unavailable or lacks PyInit__ata_core.
Preserve the backend contract and add regression coverage for the
unavailable-extension fallback path.
🪄 Autofix

Fix all unresolved CodeRabbit comments on this PR:

  • Push a commit to this branch (recommended)
  • Create a new PR with the fixes

ℹ️ Review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro Plus

Run ID: d47546d2-7780-49f2-bf1c-b079d4751713

📥 Commits

Reviewing files that changed from the base of the PR and between 093df33 and baef699.

📒 Files selected for processing (7)
  • crates/fast-mlsirm-py/src/ata_bindings.rs
  • crates/fast-mlsirm-py/src/entrypoint.rs
  • docs/changelog.d/805-ata-target-gain-rust.md
  • python/fast_mlsirm/_ata_core_loader.py
  • python/fast_mlsirm/ata.py
  • tests/test_ata.py
  • tests/test_modular_extension_coexistence.py

Comment thread crates/fast-mlsirm-py/src/ata_bindings.rs Outdated
Comment thread crates/fast-mlsirm-py/src/ata_bindings.rs Outdated
Comment thread python/fast_mlsirm/ata.py
@seonghobae
seonghobae enabled auto-merge (squash) August 12, 2026 11:59

Copy link
Copy Markdown
Contributor Author

@cwl-noema-review

Independent exact-head review requested for a57ae130685fbe4d98d2e70e15d80f40691e38c4 against current protected main 5231438df8e02bd71d88e692aaaa09a1d6c8cdc3.

Review the bounded ATA Rust-ownership slice only. Verify public target-gain arithmetic is Rust/PyO3-owned, Python retains only validation/orchestration/tie-breaking, missing native capability fails closed rather than selecting NumPy, candidate/matrix dimensions and allocations are bounded before conversion, exact errors do not leak PyO3 TypeError, and existing ATA semantics remain unchanged. Exact-head CI, fuzz, Security Scan, SAST and CodeRabbit status are successful; all review threads are resolved. Read-only review; do not modify or merge.

Copy link
Copy Markdown
Contributor Author

@opencode-agent

Perform a read-only exact-head review of a57ae130685fbe4d98d2e70e15d80f40691e38c4 against protected main 5231438df8e02bd71d88e692aaaa09a1d6c8cdc3. Review only ATA Rust ownership, bounded PyO3 input conversion, fail-closed native capability, parity and unchanged assembly semantics. Exact-head CI/fuzz/security/SAST are successful and every inline thread is resolved. Do not modify or merge.

@seonghobae
seonghobae merged commit d34a473 into main Aug 12, 2026
36 checks passed
@seonghobae
seonghobae deleted the fix/ata-target-gain-rust branch August 12, 2026 19:57
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant