refactor(inference): own observed information and second-order tests in Rust - #758
Conversation
|
Important Review skippedAuto incremental reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Pro Plus Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
📝 WalkthroughWalkthrough
ChangesObserved-information Rust ownership
Estimated code review effort: 3 (Moderate) | ~25 minutes Sequence Diagram(s)sequenceDiagram
participant PythonInference
participant PyO3Core
participant RustInference
PythonInference->>PythonInference: Collect objective evaluations
PythonInference->>PyO3Core: Pass finite-difference samples
PyO3Core->>RustInference: Build symmetric Hessian
RustInference-->>PyO3Core: Return Hessian
PyO3Core-->>PythonInference: Return Hessian
PythonInference->>PyO3Core: Pass Hessian and tolerance
PyO3Core->>RustInference: Compute eigenvalue diagnostics
RustInference-->>PyO3Core: Return status and eigenvalues
PyO3Core-->>PythonInference: Return diagnostic result
Possibly related issues
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 2📝 Generate docstrings 💡
⚔️ Resolve merge conflicts 💡
🧪 Generate unit tests (beta)
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. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 5
🤖 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/mlsirm-core/src/inference.rs`:
- Around line 116-124: In the Hessian construction flow, validate the
finite-difference denominators immediately after computing h2 and before
diagonal or mixed-partial division: reject h2 when it is zero or non-finite, and
reject 4.0 * h2 when it overflows or is non-finite, returning the documented
benign error. Also reject any non-finite assembled Hessian entries, and add
coverage for positive step values whose square underflows or overflows.
- Line 65: Replace or guard the jacobi_symmetric_eigen call in the inference
path so dimensions accepted by the public Python API cannot trigger unbounded
cubic diagnostic work. Prefer a scalable symmetric eigensolver; otherwise reject
oversized n before the call with a documented benign exception, and align the
Rust dimension limit with the Python API limit while preserving valid-input
behavior.
- Around line 48-49: Use checked arithmetic for all dimension-derived sizes in
vcov_from_hessian, second_order_test, finite_difference_hessian, and
standard_errors_from_vcov, including n * n, n * (n - 1) / 2, k * k, and p * p.
Return each function’s existing benign error before validation or allocation
when checked multiplication or subtraction overflows; update both the anchor at
crates/mlsirm-core/src/inference.rs:48-49 and the sibling at
crates/mlsirm-core/src/inference.rs:101-117.
In `@docs/doctoring/observed_information_rust_ownership.md`:
- Around line 10-19: Expand the documentation around the “Standards and
literature (APA 7th)” section to explain that this implementation computes
penalized-objective curvature using central finite differences, rather than
Oakes’s direct EM observed-information identity or the covariance estimators
compared by Pritikin. State the supported interpretation and document numerical
limitations caused by the chosen finite-difference step size, while retaining
the full citations and links required by the coding guidelines.
In `@python/fast_mlsirm/inference.py`:
- Around line 60-78: Remove the `np.eye(n)` dependency from the
finite-difference loop and replace it with one mutable parameter workspace
initialized from `x0`. For each evaluation in the loop around `diag_plus`,
`diag_minus`, and the off-diagonal arrays, adjust only coordinates `i` and `j`
by ±`h`, call `objective`, then reset those coordinates before the next
evaluation; preserve the existing result ordering and values without allocating
a dense identity matrix.
🪄 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: c88aa4a0-9170-44e2-a9e3-3f14da55fc80
📒 Files selected for processing (7)
CHANGELOG.mdcrates/fast-mlsirm-py/src/lib.rscrates/mlsirm-core/src/inference.rsdocs/changelog.d/734-observed-information-rust-ownership.mddocs/doctoring/observed_information_rust_ownership.mdpython/fast_mlsirm/inference.pytests/test_observed_information_rust_ownership.py
| if n == 0 || hessian.len() != n * n { | ||
| return Err("hessian must be a square matrix".into()); |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
ast-grep outline crates/mlsirm-core/src/inference.rs --items all \
--match 'second_order_test|finite_difference_hessian'
rg -n -C 3 'n \* n|n \* \(n - 1\)|checked_mul|checked_sub' \
crates/mlsirm-core/src/inference.rsRepository: ContextualWisdomLab/fast-mlsirm
Length of output: 1786
🏁 Script executed:
#!/bin/bash
set -euo pipefail
ast-grep outline crates/mlsirm-core/src/inference.rs --items all
sed -n '1,190p' crates/mlsirm-core/src/inference.rs
rg -n -C 4 'vcov_from_hessian|second_order_test|finite_difference_hessian|standard_errors_from_vcov' crates testsRepository: ContextualWisdomLab/fast-mlsirm
Length of output: 44373
🏁 Script executed:
#!/bin/bash
set -euo pipefail
sed -n '155,289p' crates/mlsirm-core/src/inference.rs
python3 - <<'PY'
U = 2**64
MASK = U - 1
def release_mul(a, b):
return (a * b) & MASK
def release_sub(a, b):
return (a - b) & MASK
for n in [2**32, 2**32 + 1, 2**63, 2**64 - 1]:
square = release_mul(n, n)
off_n = release_mul(n, release_sub(n, 1)) // 2
print(f"n={n}: n*n wraps to {square}, n*(n-1)/2 evaluates to {off_n}")
print("debug overflow thresholds:")
print("largest n with n*n <= usize::MAX:", 2**32 - 1)
print("largest n with n*(n-1) <= usize::MAX:", 2**32)
PYRepository: ContextualWisdomLab/fast-mlsirm
Length of output: 4633
Use checked arithmetic for every dimension-derived size.
If n * n or n * (n - 1) / 2 overflows, return the function's existing benign error before validation or allocation. Apply this to vcov_from_hessian, second_order_test, finite_difference_hessian, and standard_errors_from_vcov. This must also protect the helper k * k and p * p allocations.
📍 Affects 1 file
crates/mlsirm-core/src/inference.rs#L48-L49(this comment)crates/mlsirm-core/src/inference.rs#L101-L117
🤖 Prompt for 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.
In `@crates/mlsirm-core/src/inference.rs` around lines 48 - 49, Use checked
arithmetic for all dimension-derived sizes in vcov_from_hessian,
second_order_test, finite_difference_hessian, and standard_errors_from_vcov,
including n * n, n * (n - 1) / 2, k * k, and p * p. Return each function’s
existing benign error before validation or allocation when checked
multiplication or subtraction overflows; update both the anchor at
crates/mlsirm-core/src/inference.rs:48-49 and the sibling at
crates/mlsirm-core/src/inference.rs:101-117.
Source: Coding guidelines
| symmetric[j * n + i] = mean; | ||
| } | ||
| } | ||
| let (mut evals, _) = jacobi_symmetric_eigen(&symmetric, n)?; |
There was a problem hiding this comment.
🚀 Performance & Scalability | 🟠 Major | 🏗️ Heavy lift
Bound or replace the dense Jacobi diagnostic.
jacobi_symmetric_eigen uses up to 64 sweeps of cubic work. The public Python API admits n == 5_000. That input can require about 12 trillion coordinate updates before this call returns or reports non-convergence.
Use a scalable symmetric eigensolver, or reject dimensions that the fixed-sweep Jacobi method cannot serve before calling it. Keep the Python and Rust limits aligned. As per coding guidelines, “For fuzzed arbitrary input, code must either succeed or raise a documented benign exception; panics, hangs, AssertionError, KeyError, IndexError, and TypeError are bugs.”
🤖 Prompt for 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.
In `@crates/mlsirm-core/src/inference.rs` at line 65, Replace or guard the
jacobi_symmetric_eigen call in the inference path so dimensions accepted by the
public Python API cannot trigger unbounded cubic diagnostic work. Prefer a
scalable symmetric eigensolver; otherwise reject oversized n before the call
with a documented benign exception, and align the Rust dimension limit with the
Python API limit while preserving valid-input behavior.
Source: Coding guidelines
| let h2 = step * step; | ||
| let mut hessian = vec![0.0_f64; n * n]; | ||
| for i in 0..n { | ||
| hessian[i * n + i] = (diag_plus[i] - 2.0 * base + diag_minus[i]) / h2; | ||
| } | ||
| let mut k = 0usize; | ||
| for i in 0..n { | ||
| for j in (i + 1)..n { | ||
| let value = (off_pp[k] - off_pm[k] - off_mp[k] + off_mm[k]) / (4.0 * h2); |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Reject unusable finite-difference denominators.
A finite positive step can still make step * step equal 0.0 or inf. Also, 4.0 * h2 can overflow. The current code then returns zero, infinite, or NaN Hessian entries instead of a documented error.
Validate h2 and the mixed-partial denominator before division. Reject non-finite assembled entries. Add tests for underflow and overflow step values. As per coding guidelines, “For fuzzed arbitrary input, code must either succeed or raise a documented benign exception; panics, hangs, AssertionError, KeyError, IndexError, and TypeError are bugs.”
🤖 Prompt for 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.
In `@crates/mlsirm-core/src/inference.rs` around lines 116 - 124, In the Hessian
construction flow, validate the finite-difference denominators immediately after
computing h2 and before diagonal or mixed-partial division: reject h2 when it is
zero or non-finite, and reject 4.0 * h2 when it overflows or is non-finite,
returning the documented benign error. Also reject any non-finite assembled
Hessian entries, and add coverage for positive step values whose square
underflows or overflows.
Source: Coding guidelines
| ## Standards and literature (APA 7th) | ||
|
|
||
| Pritikin, J. N. (2017). A comparison of parameter covariance estimation methods | ||
| for item response models in an expectation-maximization framework. *Cogent | ||
| Psychology, 4*(1), Article 1279435. | ||
| https://doi.org/10.1080/23311908.2017.1279435 | ||
|
|
||
| Oakes, D. (1999). Direct calculation of the information matrix via the EM | ||
| algorithm. *Journal of the Royal Statistical Society Series B: Statistical | ||
| Methodology, 61*(2), 479–482. https://doi.org/10.1111/1467-9868.00188 |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick win
Add the required literature basis and scope statement.
The citations alone do not explain the supported method. Oakes derives a direct observed-information identity for EM. Pritikin compares covariance estimators for EM-based IRT models. This implementation instead uses central finite differences of the penalized objective. State that distinction, the supported interpretation, and the numerical limitations of step-size-based curvature diagnostics. (ideas.repec.org)
As per coding guidelines, “For substantive feature or process PRs, include permissible research PDFs with full citations, or cite, link, and summarize them when redistribution is not permissible.”
🤖 Prompt for 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.
In `@docs/doctoring/observed_information_rust_ownership.md` around lines 10 - 19,
Expand the documentation around the “Standards and literature (APA 7th)” section
to explain that this implementation computes penalized-objective curvature using
central finite differences, rather than Oakes’s direct EM observed-information
identity or the covariance estimators compared by Pritikin. State the supported
interpretation and document numerical limitations caused by the chosen
finite-difference step size, while retaining the full citations and links
required by the coding guidelines.
Source: Coding guidelines
| # Python evaluates the scalar objective at FD offsets; Rust owns the | ||
| # finite-difference coefficients and symmetrised matrix assembly. | ||
| diag_plus = np.empty(n, dtype=np.float64) | ||
| diag_minus = np.empty(n, dtype=np.float64) | ||
| off_n = n * (n - 1) // 2 | ||
| off_pp = np.empty(off_n, dtype=np.float64) | ||
| off_pm = np.empty(off_n, dtype=np.float64) | ||
| off_mp = np.empty(off_n, dtype=np.float64) | ||
| off_mm = np.empty(off_n, dtype=np.float64) | ||
| k = 0 | ||
| for i in range(n): | ||
| x_plus = x0 + h * eye[i] | ||
| x_minus = x0 - h * eye[i] | ||
| hessian[i, i] = (objective(x_plus) - 2.0 * base + objective(x_minus)) / (h * h) | ||
| diag_plus[i] = objective(x0 + h * eye[i]) | ||
| diag_minus[i] = objective(x0 - h * eye[i]) | ||
| for j in range(i + 1, n): | ||
| f_pp = objective(x0 + h * eye[i] + h * eye[j]) | ||
| f_pm = objective(x0 + h * eye[i] - h * eye[j]) | ||
| f_mp = objective(x0 - h * eye[i] + h * eye[j]) | ||
| f_mm = objective(x0 - h * eye[i] - h * eye[j]) | ||
| value = (f_pp - f_pm - f_mp + f_mm) / (4.0 * h * h) | ||
| hessian[i, j] = value | ||
| hessian[j, i] = value | ||
| off_pp[k] = objective(x0 + h * eye[i] + h * eye[j]) | ||
| off_pm[k] = objective(x0 + h * eye[i] - h * eye[j]) | ||
| off_mp[k] = objective(x0 - h * eye[i] + h * eye[j]) | ||
| off_mm[k] = objective(x0 - h * eye[i] - h * eye[j]) | ||
| k += 1 |
There was a problem hiding this comment.
🚀 Performance & Scalability | 🟠 Major | ⚡ Quick win
Remove the dense identity workspace.
The offset loop depends on np.eye(n) from Line 57. At the admitted 5,000-parameter limit, that workspace adds 200 MB to an already dense diagnostic path.
Use one mutable parameter workspace. Adjust and reset only coordinates i and j for each objective evaluation. As per coding guidelines, “Avoid unnecessary intermediate NumPy allocations and prefer einsum or BLAS-backed forms where appropriate.”
🤖 Prompt for 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.
In `@python/fast_mlsirm/inference.py` around lines 60 - 78, Remove the `np.eye(n)`
dependency from the finite-difference loop and replace it with one mutable
parameter workspace initialized from `x0`. For each evaluation in the loop
around `diag_plus`, `diag_minus`, and the off-diagonal arrays, adjust only
coordinates `i` and `j` by ±`h`, call `objective`, then reset those coordinates
before the next evaluation; preserve the existing result ordering and values
without allocating a dense identity matrix.
Source: Coding guidelines
…in Rust Public observed_information assembles finite-difference Hessians in the Rust core from evaluated objective samples; second_order_test eigenvalue diagnostics are Rust-owned. Ownership sentinels, quadratic recovery unit tests, changelog, and APA doctoring included. Supersedes draft #734 once green.
671290f to
207e259
Compare
1 participant
Why
Draft #734 is a RED ownership contract (test-only). Production Hessian assembly and second-order diagnostics still ran fully in Python.
What
finite_difference_hessianowns FD coefficients + symmetrisationsecond_order_testowns Jacobi eigenvalue diagnosticsobserved_information/second_order_testVerification
cargo test -p mlsirm-core --lib inference::tests(7 pass)pytest tests/test_observed_information_rust_ownership.py+ related suite (24 pass)Supersedes #734 once product gates are green.
Summary by CodeRabbit
New Features
Documentation
Tests