feat(index, random): implement index_bool, index_take, and Rng struct - #263
Open
rajesh-puripanda wants to merge 2 commits into
Open
feat(index, random): implement index_bool, index_take, and Rng struct#263rajesh-puripanda wants to merge 2 commits into
rajesh-puripanda wants to merge 2 commits into
Conversation
Add boolean mask indexing (index_bool) and integer array indexing along an axis (index_take) to the mohu-index crate. index_bool filters a buffer by a boolean mask, returning a 1D buffer of selected elements. Supports all dtypes via dispatch_dtype. index_take selects elements along a given axis using an array of integer indices, with support for negative indices and bounds checking. Uses dispatch_dtype for runtime type dispatch. Both functions are fully safe (no unsafe code) and include comprehensive tests. Signed-off-by: rajesh-puripanda <rajesh.puripanda@icloud.com>
5 tasks
PR Check Summary
CI will run: build, test, clippy, fmt, cargo-deny, DCO, semver. |
📝 WalkthroughWalkthroughThis PR implements two advanced indexing functions for mohu-index: ChangesAdvanced Indexing API
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 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 |
Add convenience Rng struct wrapping Pcg64 with three sampling methods: - uniform: f64 values in [low, high) - normal: Gaussian variates via Box-Muller (all dtypes via dispatch) - integers: i64 values in [low, high) Includes reproducibility guarantee, domain error validation, and comprehensive tests for all methods. Signed-off-by: rajesh-puripanda <rajesh.puripanda@icloud.com>
rajesh-puripanda
changed the title
6 tasks
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 4
🧹 Nitpick comments (1)
crates/mohu-index/src/boolean.rs (1)
30-50: ⚖️ Poor tradeoffEliminate per-true
Vec<usize>allocations incrates/mohu-index/src/boolean.rs
true_coords: Vec<Vec<usize>>pluscoord.to_vec()heap-allocates oneVecper selected element.mohu-buffer::Bufferexposesget/set(&[usize])(no linear/offsetget/set), so the simplest allocation-free refactor is a two-pass walk: first counttrueelements to sizeout, then re-iterate the indices and copy matchingsrcvalues intooutdirectly (without storing coordinates).🤖 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/mohu-index/src/boolean.rs` around lines 30 - 50, The current implementation builds true_coords: Vec<Vec<usize>> by pushing coord.to_vec() for every true mask which allocates per-element; instead, do a two-pass iteration over NdIndexIter::new(src.shape()): first pass just count mask.get::<bool>(&coord)? to compute out_len and allocate out via Buffer::zeros(src.dtype(), &[out_len])?, return early if zero; second pass iterate again and maintain a single usize write_idx counter that increments for each true mask and directly call src.get::<T>(&coord)? and out.set::<T>(&[write_idx], val)? (adjust the copy_bool! macro to use the write_idx counter instead of indexing into true_coords). Remove true_coords and coord.to_vec() usage entirely.
🤖 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/mohu-index/src/boolean.rs`:
- Around line 82-88: The test test_index_bool_all_true currently builds a
Buffer<f64> and compares slices with assert_eq, but
mohu_testing::approx::assert_allclose only accepts &[f32]; change the test to
use Buffer::from_slice::<f32> for src, produce the mask as before, call
index_bool(&src, &mask).unwrap(), then call
mohu_testing::approx::assert_allclose on result.as_slice::<f32>().unwrap() with
the expected &[1.0_f32, 2.0_f32, 3.0_f32]; keep the test name and use the same
Buffer and index_bool symbols so only the element type and assertion change.
- Around line 17-28: Replace the manual early-returns in the function
index_bool: instead of the explicit if checks that call return Err(...), use the
mohu-error ensure! macro (ensure!(mask.dtype() == DType::Bool,
MohuError::DomainError { ... }) and ensure!(mask.shape() == src.shape(),
MohuError::BoolIndexShapeMismatch { ... })), preserving the same error
constructors; also update the test test_index_bool_all_true to use f32 data and
compare using moh u_testing::approx::assert_allclose (switch the slice type from
f64 to f32 and call assert_allclose on the resulting &[f32]) so the approximate
comparison helper is used.
In `@crates/mohu-index/src/take.rs`:
- Around line 67-71: The branch that handles empty-source currently sets
out_shape[axis] = 0 which is incorrect when the source is empty due to a
non-axis dimension; instead set out_shape[axis] = num_indices so the output
shape matches NumPy take semantics (works for both src.len() == 0 and
num_indices == 0). Update the code around the if checking num_indices == 0 ||
src.len() == 0 in take.rs to assign out_shape[axis] = num_indices before
returning Buffer::zeros(src.dtype(), &out_shape).
- Around line 17-36: The error handling in index_take needs to use the project's
ensure!/bail! macros and construct AxisOutOfRange with its required valid
string: replace the explicit returns of Err(MohuError::TooManyIndices { .. }),
Err(MohuError::AxisOutOfRange { .. }), and Err(MohuError::DomainError { .. }) in
the function index_take with ensure! or bail! calls (e.g., ensure!(ndim != 0,
TooManyIndices{...}) or bail!(DomainError{...}) as appropriate), and when
constructing MohuError::AxisOutOfRange include the missing valid: String field
(for example "0..ndim") so the variant is fully initialized; also verify the
empty-input/output shape handling remains correct and update any tests that
compare floats to use mohu_testing::approx::assert_allclose instead of exact
equality.
---
Nitpick comments:
In `@crates/mohu-index/src/boolean.rs`:
- Around line 30-50: The current implementation builds true_coords:
Vec<Vec<usize>> by pushing coord.to_vec() for every true mask which allocates
per-element; instead, do a two-pass iteration over
NdIndexIter::new(src.shape()): first pass just count mask.get::<bool>(&coord)?
to compute out_len and allocate out via Buffer::zeros(src.dtype(), &[out_len])?,
return early if zero; second pass iterate again and maintain a single usize
write_idx counter that increments for each true mask and directly call
src.get::<T>(&coord)? and out.set::<T>(&[write_idx], val)? (adjust the
copy_bool! macro to use the write_idx counter instead of indexing into
true_coords). Remove true_coords and coord.to_vec() usage entirely.
🪄 Autofix (Beta)
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: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 1bfdb847-85eb-46f6-b69e-dc10e3590019
📒 Files selected for processing (2)
crates/mohu-index/src/boolean.rscrates/mohu-index/src/take.rs
Comment on lines
+17
to
+36
| pub fn index_take(src: &Buffer, indices: &Buffer, axis: usize) -> MohuResult<Buffer> { | ||
| let ndim = src.ndim(); | ||
| if ndim == 0 { | ||
| return Err(MohuError::TooManyIndices { | ||
| given: 1, | ||
| ndim: 0, | ||
| }); | ||
| } | ||
| if axis >= ndim { | ||
| return Err(MohuError::AxisOutOfRange { | ||
| axis: axis as i64, | ||
| ndim, | ||
| }); | ||
| } | ||
| if indices.ndim() != 1 || indices.dtype() != DType::I64 { | ||
| return Err(MohuError::DomainError { | ||
| op: "index_take", | ||
| reason: "indices must be a 1D I64 array".into(), | ||
| }); | ||
| } |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
ast-grep --pattern 'pub enum MohuError { $$$ }'
rg -nP '\b(FancyIndexOutOfBounds|TooManyIndices|AxisOutOfRange|IndexOutOfBounds)\b' --type=rust -g '*error*' -C1Repository: mohu-org/mohu
Length of output: 37948
Fix index_take error construction + use ensure!/bail! for guards
- In
crates/mohu-index/src/take.rs,MohuError::AxisOutOfRangemust include the requiredvalid: Stringfield (the currentAxisOutOfRange { axis: ..., ndim: ... }construction omits it). - Replace the guard/out-of-bounds
return Err(MohuError::...)paths withensure!/bail!per the project error-handling rules. - Double-check the empty-input/output shape handling and update any float equality assertions in tests to use
mohu_testing::approx::assert_allclose.
🤖 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/mohu-index/src/take.rs` around lines 17 - 36, The error handling in
index_take needs to use the project's ensure!/bail! macros and construct
AxisOutOfRange with its required valid string: replace the explicit returns of
Err(MohuError::TooManyIndices { .. }), Err(MohuError::AxisOutOfRange { .. }),
and Err(MohuError::DomainError { .. }) in the function index_take with ensure!
or bail! calls (e.g., ensure!(ndim != 0, TooManyIndices{...}) or
bail!(DomainError{...}) as appropriate), and when constructing
MohuError::AxisOutOfRange include the missing valid: String field (for example
"0..ndim") so the variant is fully initialized; also verify the
empty-input/output shape handling remains correct and update any tests that
compare floats to use mohu_testing::approx::assert_allclose instead of exact
equality.
Comment on lines
+67
to
+71
| if num_indices == 0 || src.len() == 0 { | ||
| let mut out_shape: Vec<usize> = src.shape().to_vec(); | ||
| out_shape[axis] = 0; | ||
| return Buffer::zeros(src.dtype(), &out_shape); | ||
| } |
There was a problem hiding this comment.
Empty-source path produces wrong output shape along the axis.
When src.len() == 0 because a non-axis dimension is zero (e.g. shape [0, 3], axis = 1, indices = [0, 1]), the bounds checks pass but this branch forces out_shape[axis] = 0, yielding [0, 0] instead of the expected [0, num_indices] (NumPy take semantics). Using num_indices also stays correct for the num_indices == 0 case.
🐛 Proposed fix
if num_indices == 0 || src.len() == 0 {
let mut out_shape: Vec<usize> = src.shape().to_vec();
- out_shape[axis] = 0;
+ out_shape[axis] = num_indices;
return Buffer::zeros(src.dtype(), &out_shape);
}📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| if num_indices == 0 || src.len() == 0 { | |
| let mut out_shape: Vec<usize> = src.shape().to_vec(); | |
| out_shape[axis] = 0; | |
| return Buffer::zeros(src.dtype(), &out_shape); | |
| } | |
| if num_indices == 0 || src.len() == 0 { | |
| let mut out_shape: Vec<usize> = src.shape().to_vec(); | |
| out_shape[axis] = num_indices; | |
| return Buffer::zeros(src.dtype(), &out_shape); | |
| } |
🤖 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/mohu-index/src/take.rs` around lines 67 - 71, The branch that handles
empty-source currently sets out_shape[axis] = 0 which is incorrect when the
source is empty due to a non-axis dimension; instead set out_shape[axis] =
num_indices so the output shape matches NumPy take semantics (works for both
src.len() == 0 and num_indices == 0). Update the code around the if checking
num_indices == 0 || src.len() == 0 in take.rs to assign out_shape[axis] =
num_indices before returning Buffer::zeros(src.dtype(), &out_shape).
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Implements two issues:
#149 — Advanced indexing
\index_bool(src, mask) -> Buffer\
Boolean mask filtering — returns 1D buffer of elements where mask is true.
\index_take(src, indices, axis) -> Buffer\
Integer array indexing along an axis (NumPy's \ ake). Supports negative indices and bounds checking.
#140 — Random number generator
\Rng\ struct
ew(seed: u64)\ — seeded PCG-64-DXSM engine
ormal(shape, mean, std)\ — f64 Gaussian via Box-Muller
Both use safe Rust only, with \dispatch_dtype!\ for runtime type dispatch (index) and direct engine usage (Rng). Includes comprehensive tests for edge cases, reproducibility, and error handling.
Closes #149
Closes #140