⚡ Bolt: Optimize grad_alpha computation via dense matrix multiplication - #202
⚡ Bolt: Optimize grad_alpha computation via dense matrix multiplication#202seonghobae wants to merge 2 commits into
Conversation
|
👋 Jules, reporting for duty! I'm here to lend a hand with this pull request. When you start a review, I'll add a 👀 emoji to each comment to let you know I've read it. I'll focus on feedback directed at me and will do my best to stay out of conversations between you and other bots or reviewers to keep the noise down. I'll push a commit with your requested changes shortly after. Please note there might be a delay between these steps, but rest assured I'm on the job! For more direct control, you can switch me to Reactive Mode. When this mode is on, I will only act on comments where you specifically mention me with New to Jules? Learn more at jules.google/docs. For security, I will only act on instructions from the user who triggered this task. |
There was a problem hiding this comment.
Pull request overview
Note
Copilot couldn't run its full agentic review because it didn't start before the timeout. Make sure your repository has a runner available, or add a copilot-code-review.yml file specifying one with the runs-on attribute. See the docs for more details.
Optimizes grad_alpha computation in the Python objective by replacing an element-wise multiply + reduction with a BLAS-backed matrix multiplication and targeted indexing to reduce large intermediate allocations.
Changes:
- Replaced
(e * theta[:, factors]).sum(axis=0)with(e.T @ theta)[rows, factors]forgrad_alpha. - Added an internal optimization note in
.jules/bolt.mddocumenting the pattern.
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 3 comments.
| File | Description |
|---|---|
| python/fast_mlsirm/objective.py | Updates grad_alpha gradient aggregation to use matmul + indexing to reduce N×J intermediates. |
| .jules/bolt.md | Documents the optimization pattern and rationale for future reference. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| # Optimized gradient computation: avoid intermediate array allocation by using matrix multiplication | ||
| grad_alpha = (e.T @ params.theta)[np.arange(e.shape[1]), factors] * a |
| grad_alpha = np.zeros_like(params.alpha) | ||
| if free_alpha: | ||
| grad_alpha = (e * params.theta[:, factors]).sum(axis=0) * a | ||
| # Optimized gradient computation: avoid intermediate array allocation by using matrix multiplication |
| **Learning:** During gradient calculation, `float((e * (-gamma * distance)).sum())` creates two full-size `(N, J)` arrays: one for the scaled distance and one for the element-wise multiplication before reduction. | ||
| **Action:** Replace `(A * B).sum()` with `np.vdot(A, B)` when scalar reduction is needed over matrix multiplication (where `B` can incorporate scalars naturally like `-gamma * np.vdot(A, B)`). This entirely avoids the 2D array allocation overhead and yields order-of-magnitude improvements in scalar gradient components. | ||
| ## 2025-05-19 - Matrix Multiplication for Factor Gradient Calculation | ||
| **Learning:** Calculating gradients for factor arrays using boolean element-wise extraction and reduction like `(e * theta[:, factors]).sum(axis=0)` loops through columns and allocates large intermediate arrays (N x J). For large N (e.g., 5000), this significantly degrades performance. |
…nction. * Implemented safe fallback to `np.hypot` for spatial distances exceeding limits. * Added `math.isfinite` check for `eps_distance`. * Assured non-fractional validation on `factor_id`. * Safe-guarded boolean masking by applying `.copy()` to ensure mutability doesn't affect callers.
| factors = np.asarray(factor_id, dtype=np.float64) | ||
| if not np.all(factors == np.floor(factors)): | ||
| raise ValueError("factor_id must contain integer values") | ||
| factors = factors.astype(np.int64) |
| if max_val > 1e100: | ||
| diff = params.xi[:, None, :] - params.zeta[None, :, :] | ||
| dist_sq = np.zeros(diff.shape[:2], dtype=diff.dtype) | ||
| for i in range(diff.shape[-1]): | ||
| dist_sq = np.hypot(dist_sq, diff[..., i]) | ||
| distance = np.sqrt(dist_sq**2 + eps_distance) | ||
| else: | ||
| xi_sq = np.einsum('ij,ij->i', params.xi, params.xi) | ||
| zeta_sq = np.einsum('ij,ij->i', params.zeta, params.zeta) | ||
| dist_sq = xi_sq[:, None] + zeta_sq[None, :] - 2 * np.dot(params.xi, params.zeta.T) | ||
| dist_sq = np.maximum(dist_sq, 0.0) | ||
| distance = np.sqrt(dist_sq + eps_distance) |
| factors = np.asarray(factor_id, dtype=np.float64) | ||
| if not np.all(factors == np.floor(factors)): | ||
| raise ValueError("factor_id must contain integer values") |
|
중복/번들 정리: 배경: 2026-07-14 이후 조직 coverage-evidence 인프라 문제로 모든 PR이 REQUEST_CHANGES 상태였습니다(인프라 수정: ContextualWisdomLab/.github#611). 필요 시 재오픈 가능합니다. Generated by Claude Code |
Understood. Acknowledging that this work is now obsolete and stopping work on this task. (이해가 완료되었습니다. 현재 PR이 더 이상 필요하지 않음을 확인하였으며, 이 작업에 대한 추가 수정을 중단하겠습니다.) |
💡 What: Replaced the loop-based array extraction and element-wise reduction
(e * params.theta[:, factors]).sum(axis=0)with optimized dense matrix multiplication(e.T @ params.theta)[np.arange(e.shape[1]), factors].🎯 Why: To prevent large, unnecessary
N x Jintermediate array allocations during gradient aggregation over boolean subsets, reducing both memory pressure and execution time.📊 Impact: Reduces matrix multiplication overhead by avoiding Python looping over factors, leveraging pure C/BLAS computation for dense matrix arrays.
🔬 Measurement: Python scripts verified matching calculations.
uv run pytest tests/andcargo testdemonstrated stable numeric parity and fast execution.PR created automatically by Jules for task 12308938969966372113 started by @seonghobae