Density Preserving UMAP - #7860
Closed
viclafargue wants to merge 6 commits into
Closed
Conversation
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
viclafargue
added
improvement
cjnolet
requested changes
Mar 10, 2026
|
Check out this pull request on See visual diffs & provide feedback on Jupyter Notebooks. Powered by ReviewNB |
viclafargue
force-pushed
the
density-preserving-umap
branch
from
March 20, 2026 17:04
bb8903c to
c0229f6
Compare
viclafargue
marked this pull request as ready for review
March 20, 2026 17:05
Contributor
Author
I could remove 3 custom kernels to replace them with RAFT primitives. |
📝 WalkthroughSummary by CodeRabbit
WalkthroughAdds DensMAP support: new UMAPParams fields, GPU kernels for edge distances and DensMAP precompute/epoch stats, integrates DensMAP into optimization kernels and runners, exposes parameters in Python bindings, and adds tests for densMAP behavior and validation. Changes
Estimated code review effort🎯 4 (Complex) | ⏱️ ~65 minutes Possibly related PRs
Suggested reviewers
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
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: 6
🧹 Nitpick comments (4)
python/cuml/tests/test_umap_densmap.py (2)
31-42: Consider documenting the trustworthiness threshold choice.The threshold of
>= 0.90is reasonable for iris, but adding a brief comment explaining the expected range would help future maintainers understand if a slightly lower value is acceptable or indicates a regression.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@python/cuml/tests/test_umap_densmap.py` around lines 31 - 42, Add a short comment in test_densmap_trustworthiness_on_iris explaining why the trustworthiness threshold is set to >= 0.90 (e.g., typical trustworthiness range for iris with densmap/cuUMAP, acceptable variance, and that values slightly below may indicate regression); update the test near the cuUMAP(...).fit_transform call and the subsequent trustworthiness(iris.data, embedding, n_neighbors=10) assertion to document the expected range and rationale.
59-81: Consider widening the reproducibility tolerance or adding retries.The
mean_diff < 1.0threshold accounts for atomicAdd non-determinism, but this test could still be flaky on certain GPU architectures or under high system load. Consider either widening the threshold or marking as@pytest.mark.flaky(reruns=2)if the test framework supports it.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@python/cuml/tests/test_umap_densmap.py` around lines 59 - 81, The test_densmap_reproducibility test is flaky because mean_diff < 1.0 can fail on some GPUs; either relax the tolerance (e.g., change the assertion to mean_diff < 1.5 or another empirically-determined value) or add a retry marker (e.g., annotate the test function with `@pytest.mark.flaky`(reruns=2) or your test-suite's equivalent) to reduce spurious failures; update the assertion or add the marker on the test_densmap_reproducibility function that calls cuUMAP(...).fit_transform to implement the chosen fix.cpp/src/umap/runner.cuh (1)
351-357: Verify sparse input handling with densmap=True.The
if constexprguard prevents compilation errors for sparse inputs, but ifparams->densmapis true with sparse data,dmremainsnullptrand densMAP silently degrades to standard UMAP. The Python layer raisesNotImplementedErrorfor sparse+densmap, but consider adding a runtime check here for defense-in-depth.Optional: Add runtime assertion for sparse inputs with densmap
std::unique_ptr<DensMap::DensMapData<value_t>> dm; if constexpr (detail::has_dense_X<umap_inputs>::value) { if (params->densmap) { dm = DensMap::densmap_precompute<value_t, nnz_t, TPB_X>( inputs.X, inputs.n, inputs.d, graph, params, stream); } + } else { + RAFT_EXPECTS(!params->densmap, "densMAP requires dense input data"); }🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@cpp/src/umap/runner.cuh` around lines 351 - 357, Add a runtime guard inside the densmap precompute block so that if params->densmap is true but the compile-time dense-input trait (detail::has_dense_X<umap_inputs>::value) is false (i.e., sparse inputs like inputs.X), the code raises a clear error instead of silently leaving dm null; update the check near the dm allocation/usage (where dm and DensMap::densmap_precompute are referenced) to throw a descriptive std::runtime_error or use an assert when params->densmap && !has_dense_X, so callers receive an explicit failure for sparse+densmap.cpp/src/umap/edge_metric.cuh (1)
17-31: Consider using RAFT's warp reduction primitives.RAFT provides
raft::shfl_reduceand similar utilities in<raft/util/cuda_utils.cuh>that handle warp reductions. Using those would reduce maintenance burden and ensure consistency with the rest of the codebase.That said, these implementations are correct and work well for the current use case.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@cpp/src/umap/edge_metric.cuh` around lines 17 - 31, Replace the hand-rolled warp reductions warp_reduce_sum and warp_reduce_max with RAFT's warp reduction utilities: include <raft/util/cuda_utils.cuh> and call the appropriate raft::shfl_reduce (or equivalent raft::shfl_sum/raft::shfl_max helpers) in place of the manual __shfl_down_sync loops so the code uses raft's tested primitives for both warp_reduce_sum and warp_reduce_max; keep the same template signatures and return types so callers of warp_reduce_sum and warp_reduce_max remain unchanged.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@cpp/src/umap/densmap.cuh`:
- Around line 166-167: The code casts ro_var to float before calling std::sqrt
which can lose precision when the template parameter T is double; change the
sqrt call to use T precision (e.g., T ro_std = std::sqrt(static_cast<T>(ro_var))
or std::sqrt(ro_var)) and make the epsilon comparison use T(1e-10) (keep the
existing T(1e-10) literal) so ro_std is computed and compared at the same
precision as T; update the initialization of ro_std and the conditional that
sets it to T(1e-10) to operate on T values.
- Around line 279-286: The normalization step using raft::linalg::binaryOp that
computes log(T(1e-8) + r / p) can produce Inf when phi_sum (p) is zero for
isolated vertices; update the lambda in the binaryOp that operates on re_sum and
phi_sum (used with re_sum, phi_sum, n_vertices, stream) to protect the
denominator by adding or max-clamping with a small epsilon of type T (e.g., use
r / (p + T(1e-8)) or r / max(p, T(1e-8))) before taking the log so
division-by-zero cannot occur.
- Around line 147-153: The binaryOp call computing log(1e-8 + ro_val / mu_val)
can divide by zero when mu_val is zero (isolated vertex); update the lambda
passed to raft::linalg::binaryOp (the device functor operating on dm->R, dm->R,
mu_sum.data()) to protect the denominator by using an epsilon-protected
denominator (e.g., T denom = mu_val + T(1e-8) or denom = max(mu_val, T(1e-8)))
and compute log(T(1e-8) + ro_val / denom) to avoid Inf/NaN.
- Around line 321-327: The block computing
dm.re_mean/re_std/re_cov/cov_over_var/outer_scale can divide by zero and loses
precision; update densmap.cuh to (1) stop casting h_var to float—use T for sqrt
input so dm.re_std = sqrt(static_cast<T>(h_var) + dens_var_shift); (2) guard the
covariance denominator: when computing dm.re_cov = h_dot / T(n - 1) treat n <= 1
as a special case (set dm.re_cov = T(0) or use max(1, n-1) with explicit
comment); (3) protect re_std_sq and re_std*T(n) by checking re_std_sq > eps (use
a small epsilon based on std::numeric_limits<T>::epsilon()) before dividing and
fall back to zero or a safe default for dm.cov_over_var and dm.outer_scale; and
(4) ensure any division uses safe_division-style checks (refer to dm.re_cov,
dm.cov_over_var, dm.outer_scale, h_var, dens_var_shift, and n) so no
divide-by-zero or precision loss occurs.
- Around line 347-362: The gradient computation divides by dist_squared, phi_sum
entries, and mu_edge which can be zero; to fix, add small positive guards (e.g.,
T eps = max(T(1e-12), std::numeric_limits<T>::epsilon())) and replace direct
divides with safe denominators: use denom_dist = max(dist_squared, eps) for
dphi_term, use denom_phi_sum_k = max(phi_sum[k], eps) and denom_phi_sum_j =
max(phi_sum[j], eps) for q_jk/q_kj, and use denom_mu_edge = max(mu_edge, eps)
when dividing the final value; ensure these guarded variables are used where
dist_squared, phi_sum[k]/[j], and mu_edge appear (symbols: dist_squared,
dphi_term, q_jk, q_kj, phi_sum, mu_edge, drk, drj, outer_scale).
In `@cpp/src/umap/edge_metric.cuh`:
- Around line 113-114: The LpUnexpanded finalization currently does pow(acc[0],
T(1) / p_val) which will divide by zero if p_val == 0; update the
DT::LpUnexpanded branch in edge_metric.cuh to guard p_val by clamping it to a
small positive epsilon (e.g., T eps = std::numeric_limits<T>::epsilon() or a
tiny constant) before computing T(1)/p_val, or handle p_val==0 explicitly
(return acc[0] or a defined value) so pow is never called with a
division-by-zero; change the expression using the clamped value (referencing
metric, DT::LpUnexpanded, p_val, acc, and T).
---
Nitpick comments:
In `@cpp/src/umap/edge_metric.cuh`:
- Around line 17-31: Replace the hand-rolled warp reductions warp_reduce_sum and
warp_reduce_max with RAFT's warp reduction utilities: include
<raft/util/cuda_utils.cuh> and call the appropriate raft::shfl_reduce (or
equivalent raft::shfl_sum/raft::shfl_max helpers) in place of the manual
__shfl_down_sync loops so the code uses raft's tested primitives for both
warp_reduce_sum and warp_reduce_max; keep the same template signatures and
return types so callers of warp_reduce_sum and warp_reduce_max remain unchanged.
In `@cpp/src/umap/runner.cuh`:
- Around line 351-357: Add a runtime guard inside the densmap precompute block
so that if params->densmap is true but the compile-time dense-input trait
(detail::has_dense_X<umap_inputs>::value) is false (i.e., sparse inputs like
inputs.X), the code raises a clear error instead of silently leaving dm null;
update the check near the dm allocation/usage (where dm and
DensMap::densmap_precompute are referenced) to throw a descriptive
std::runtime_error or use an assert when params->densmap && !has_dense_X, so
callers receive an explicit failure for sparse+densmap.
In `@python/cuml/tests/test_umap_densmap.py`:
- Around line 31-42: Add a short comment in test_densmap_trustworthiness_on_iris
explaining why the trustworthiness threshold is set to >= 0.90 (e.g., typical
trustworthiness range for iris with densmap/cuUMAP, acceptable variance, and
that values slightly below may indicate regression); update the test near the
cuUMAP(...).fit_transform call and the subsequent trustworthiness(iris.data,
embedding, n_neighbors=10) assertion to document the expected range and
rationale.
- Around line 59-81: The test_densmap_reproducibility test is flaky because
mean_diff < 1.0 can fail on some GPUs; either relax the tolerance (e.g., change
the assertion to mean_diff < 1.5 or another empirically-determined value) or add
a retry marker (e.g., annotate the test function with
`@pytest.mark.flaky`(reruns=2) or your test-suite's equivalent) to reduce spurious
failures; update the assertion or add the marker on the
test_densmap_reproducibility function that calls cuUMAP(...).fit_transform to
implement the chosen fix.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 53234460-0dc2-4524-a315-0cdc12e37504
📒 Files selected for processing (10)
cpp/include/cuml/manifold/umapparams.hcpp/src/umap/densmap.cuhcpp/src/umap/edge_metric.cuhcpp/src/umap/runner.cuhcpp/src/umap/simpl_set_embed/algo.cuhcpp/src/umap/simpl_set_embed/optimize_batch_kernel.cuhcpp/src/umap/simpl_set_embed/runner.cuhpython/cuml/cuml/manifold/umap/lib.pxdpython/cuml/cuml/manifold/umap/umap.pyxpython/cuml/tests/test_umap_densmap.py
Comment on lines
+147
to
+153
| raft::linalg::binaryOp( | ||
| dm->R, | ||
| dm->R, | ||
| mu_sum.data(), | ||
| n_vertices, | ||
| [] __device__(T ro_val, T mu_val) { return log(T(1e-8) + ro_val / mu_val); }, | ||
| stream); |
There was a problem hiding this comment.
Division by mu_val can produce Inf if a vertex has no incident edges.
If a vertex has zero total edge weight after trimming (isolated vertex), mu_val will be zero, causing division by zero. Consider adding epsilon protection to the denominator.
🛡️ Suggested fix
raft::linalg::binaryOp(
dm->R,
dm->R,
mu_sum.data(),
n_vertices,
- [] __device__(T ro_val, T mu_val) { return log(T(1e-8) + ro_val / mu_val); },
+ [] __device__(T ro_val, T mu_val) { return log(T(1e-8) + ro_val / (mu_val + T(1e-10))); },
stream);📝 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
| raft::linalg::binaryOp( | |
| dm->R, | |
| dm->R, | |
| mu_sum.data(), | |
| n_vertices, | |
| [] __device__(T ro_val, T mu_val) { return log(T(1e-8) + ro_val / mu_val); }, | |
| stream); | |
| raft::linalg::binaryOp( | |
| dm->R, | |
| dm->R, | |
| mu_sum.data(), | |
| n_vertices, | |
| [] __device__(T ro_val, T mu_val) { return log(T(1e-8) + ro_val / (mu_val + T(1e-10))); }, | |
| stream); |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@cpp/src/umap/densmap.cuh` around lines 147 - 153, The binaryOp call computing
log(1e-8 + ro_val / mu_val) can divide by zero when mu_val is zero (isolated
vertex); update the lambda passed to raft::linalg::binaryOp (the device functor
operating on dm->R, dm->R, mu_sum.data()) to protect the denominator by using an
epsilon-protected denominator (e.g., T denom = mu_val + T(1e-8) or denom =
max(mu_val, T(1e-8))) and compute log(T(1e-8) + ro_val / denom) to avoid
Inf/NaN.
Comment on lines
+166
to
+167
| T ro_std = std::sqrt(static_cast<float>(ro_var)); | ||
| if (ro_std < T(1e-10)) ro_std = T(1e-10); |
There was a problem hiding this comment.
Potential precision loss when T is double.
Casting ro_var to float before sqrt loses precision when the template is instantiated with double. Use std::sqrt directly on the value, or cast to T.
🔧 Suggested fix
- T ro_std = std::sqrt(static_cast<float>(ro_var));
+ T ro_std = std::sqrt(ro_var);📝 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
| T ro_std = std::sqrt(static_cast<float>(ro_var)); | |
| if (ro_std < T(1e-10)) ro_std = T(1e-10); | |
| T ro_std = std::sqrt(ro_var); | |
| if (ro_std < T(1e-10)) ro_std = T(1e-10); |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@cpp/src/umap/densmap.cuh` around lines 166 - 167, The code casts ro_var to
float before calling std::sqrt which can lose precision when the template
parameter T is double; change the sqrt call to use T precision (e.g., T ro_std =
std::sqrt(static_cast<T>(ro_var)) or std::sqrt(ro_var)) and make the epsilon
comparison use T(1e-10) (keep the existing T(1e-10) literal) so ro_std is
computed and compared at the same precision as T; update the initialization of
ro_std and the conditional that sets it to T(1e-10) to operate on T values.
Comment on lines
+279
to
+286
| // Normalize: re_sum = log(eps + re_sum / phi_sum), then exp_neg = exp(-re_sum) | ||
| raft::linalg::binaryOp( | ||
| re_sum, | ||
| re_sum, | ||
| phi_sum, | ||
| n_vertices, | ||
| [] __device__(T r, T p) { return log(T(1e-8) + r / p); }, | ||
| stream); |
There was a problem hiding this comment.
Division by phi_sum (p) can produce Inf for isolated vertices.
Same issue as in precompute: if phi_sum[v] == 0 for an isolated vertex, the division will produce Inf. Add epsilon protection to the denominator.
🛡️ Suggested fix
raft::linalg::binaryOp(
re_sum,
re_sum,
phi_sum,
n_vertices,
- [] __device__(T r, T p) { return log(T(1e-8) + r / p); },
+ [] __device__(T r, T p) { return log(T(1e-8) + r / (p + T(1e-10))); },
stream);📝 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
| // Normalize: re_sum = log(eps + re_sum / phi_sum), then exp_neg = exp(-re_sum) | |
| raft::linalg::binaryOp( | |
| re_sum, | |
| re_sum, | |
| phi_sum, | |
| n_vertices, | |
| [] __device__(T r, T p) { return log(T(1e-8) + r / p); }, | |
| stream); | |
| // Normalize: re_sum = log(eps + re_sum / phi_sum), then exp_neg = exp(-re_sum) | |
| raft::linalg::binaryOp( | |
| re_sum, | |
| re_sum, | |
| phi_sum, | |
| n_vertices, | |
| [] __device__(T r, T p) { return log(T(1e-8) + r / (p + T(1e-10))); }, | |
| stream); |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@cpp/src/umap/densmap.cuh` around lines 279 - 286, The normalization step
using raft::linalg::binaryOp that computes log(T(1e-8) + r / p) can produce Inf
when phi_sum (p) is zero for isolated vertices; update the lambda in the
binaryOp that operates on re_sum and phi_sum (used with re_sum, phi_sum,
n_vertices, stream) to protect the denominator by adding or max-clamping with a
small epsilon of type T (e.g., use r / (p + T(1e-8)) or r / max(p, T(1e-8)))
before taking the log so division-by-zero cannot occur.
Comment on lines
+321
to
+327
| dm.re_mean = h_mean; | ||
| dm.re_std = std::sqrt(static_cast<float>(h_var) + dens_var_shift); | ||
| dm.re_cov = h_dot / T(n - 1); | ||
|
|
||
| T re_std_sq = dm.re_std * dm.re_std; | ||
| dm.cov_over_var = dm.re_cov / re_std_sq; | ||
| dm.outer_scale = dm.dens_lambda * dm.mu_tot / (dm.re_std * T(n)); |
There was a problem hiding this comment.
Multiple potential division-by-zero risks.
Several divisions can fail:
- Line 323:
n - 1whenn == 1(single vertex) - Line 326:
re_std_sqwhen variance anddens_var_shiftare both zero - Line 327:
re_std * T(n)when either is zero
Also, line 322 has the same static_cast<float> precision issue.
🛡️ Suggested guards
- dm.re_mean = h_mean;
- dm.re_std = std::sqrt(static_cast<float>(h_var) + dens_var_shift);
- dm.re_cov = h_dot / T(n - 1);
+ dm.re_mean = h_mean;
+ dm.re_std = std::sqrt(h_var + static_cast<T>(dens_var_shift));
+ if (dm.re_std < T(1e-10)) dm.re_std = T(1e-10);
+ dm.re_cov = (n > 1) ? h_dot / T(n - 1) : T(0);
T re_std_sq = dm.re_std * dm.re_std;
dm.cov_over_var = dm.re_cov / re_std_sq;
- dm.outer_scale = dm.dens_lambda * dm.mu_tot / (dm.re_std * T(n));
+ dm.outer_scale = (n > 0) ? dm.dens_lambda * dm.mu_tot / (dm.re_std * T(n)) : T(0);🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@cpp/src/umap/densmap.cuh` around lines 321 - 327, The block computing
dm.re_mean/re_std/re_cov/cov_over_var/outer_scale can divide by zero and loses
precision; update densmap.cuh to (1) stop casting h_var to float—use T for sqrt
input so dm.re_std = sqrt(static_cast<T>(h_var) + dens_var_shift); (2) guard the
covariance denominator: when computing dm.re_cov = h_dot / T(n - 1) treat n <= 1
as a special case (set dm.re_cov = T(0) or use max(1, n-1) with explicit
comment); (3) protect re_std_sq and re_std*T(n) by checking re_std_sq > eps (use
a small epsilon based on std::numeric_limits<T>::epsilon()) before dividing and
fall back to zero or a safe default for dm.cov_over_var and dm.outer_scale; and
(4) ensure any division uses safe_division-style checks (refer to dm.re_cov,
dm.cov_over_var, dm.outer_scale, h_var, dens_var_shift, and n) so no
divide-by-zero or precision loss occurs.
Comment on lines
+347
to
+362
| T pow_db = pow(dist_squared, b); | ||
| T denom = T(1.0) + a * pow_db; | ||
| T phi = T(1.0) / denom; | ||
| T dphi_term = a * b * pow_db * phi / dist_squared; | ||
|
|
||
| T q_jk = phi / phi_sum[k]; | ||
| T q_kj = phi / phi_sum[j]; | ||
|
|
||
| T b_phi_term = T(1.0) - b * (T(1.0) - phi); | ||
| T drk = q_jk * (b_phi_term * exp_neg_re_sum[k] + dphi_term); | ||
| T drj = q_kj * (b_phi_term * exp_neg_re_sum[j] + dphi_term); | ||
|
|
||
| T weight_k = R[k] - cov_over_var * (re_sum[k] - re_mean); | ||
| T weight_j = R[j] - cov_over_var * (re_sum[j] - re_mean); | ||
|
|
||
| return outer_scale * (weight_k * drk + weight_j * drj) / mu_edge; |
There was a problem hiding this comment.
Critical: Division by dist_squared causes Inf/NaN when points coincide.
When two embedding points are at the same location (dist_squared == 0), line 350 divides by zero, producing Inf/NaN that propagates through the gradient computation and corrupts the optimization. This is likely to occur during early iterations or if points collapse.
Additionally, lines 352-353 divide by phi_sum and line 362 divides by mu_edge, which could also be zero.
🐛 Suggested fix with epsilon guards
template <typename T>
DI T compute_densmap_grad_coeff(T dist_squared,
T a,
T b,
int j,
int k,
T mu_edge,
T re_mean,
T cov_over_var,
T outer_scale,
const T* __restrict__ R,
const T* __restrict__ re_sum,
const T* __restrict__ phi_sum,
const T* __restrict__ exp_neg_re_sum)
{
+ constexpr T EPS = T(1e-10);
+ // Guard against zero distance
+ if (dist_squared < EPS) return T(0);
+
T pow_db = pow(dist_squared, b);
T denom = T(1.0) + a * pow_db;
T phi = T(1.0) / denom;
T dphi_term = a * b * pow_db * phi / dist_squared;
- T q_jk = phi / phi_sum[k];
- T q_kj = phi / phi_sum[j];
+ T q_jk = phi / (phi_sum[k] + EPS);
+ T q_kj = phi / (phi_sum[j] + EPS);
T b_phi_term = T(1.0) - b * (T(1.0) - phi);
T drk = q_jk * (b_phi_term * exp_neg_re_sum[k] + dphi_term);
T drj = q_kj * (b_phi_term * exp_neg_re_sum[j] + dphi_term);
T weight_k = R[k] - cov_over_var * (re_sum[k] - re_mean);
T weight_j = R[j] - cov_over_var * (re_sum[j] - re_mean);
- return outer_scale * (weight_k * drk + weight_j * drj) / mu_edge;
+ return outer_scale * (weight_k * drk + weight_j * drj) / (mu_edge + EPS);
}📝 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
| T pow_db = pow(dist_squared, b); | |
| T denom = T(1.0) + a * pow_db; | |
| T phi = T(1.0) / denom; | |
| T dphi_term = a * b * pow_db * phi / dist_squared; | |
| T q_jk = phi / phi_sum[k]; | |
| T q_kj = phi / phi_sum[j]; | |
| T b_phi_term = T(1.0) - b * (T(1.0) - phi); | |
| T drk = q_jk * (b_phi_term * exp_neg_re_sum[k] + dphi_term); | |
| T drj = q_kj * (b_phi_term * exp_neg_re_sum[j] + dphi_term); | |
| T weight_k = R[k] - cov_over_var * (re_sum[k] - re_mean); | |
| T weight_j = R[j] - cov_over_var * (re_sum[j] - re_mean); | |
| return outer_scale * (weight_k * drk + weight_j * drj) / mu_edge; | |
| template <typename T> | |
| DI T compute_densmap_grad_coeff(T dist_squared, | |
| T a, | |
| T b, | |
| int j, | |
| int k, | |
| T mu_edge, | |
| T re_mean, | |
| T cov_over_var, | |
| T outer_scale, | |
| const T* __restrict__ R, | |
| const T* __restrict__ re_sum, | |
| const T* __restrict__ phi_sum, | |
| const T* __restrict__ exp_neg_re_sum) | |
| { | |
| constexpr T EPS = T(1e-10); | |
| // Guard against zero distance | |
| if (dist_squared < EPS) return T(0); | |
| T pow_db = pow(dist_squared, b); | |
| T denom = T(1.0) + a * pow_db; | |
| T phi = T(1.0) / denom; | |
| T dphi_term = a * b * pow_db * phi / dist_squared; | |
| T q_jk = phi / (phi_sum[k] + EPS); | |
| T q_kj = phi / (phi_sum[j] + EPS); | |
| T b_phi_term = T(1.0) - b * (T(1.0) - phi); | |
| T drk = q_jk * (b_phi_term * exp_neg_re_sum[k] + dphi_term); | |
| T drj = q_kj * (b_phi_term * exp_neg_re_sum[j] + dphi_term); | |
| T weight_k = R[k] - cov_over_var * (re_sum[k] - re_mean); | |
| T weight_j = R[j] - cov_over_var * (re_sum[j] - re_mean); | |
| return outer_scale * (weight_k * drk + weight_j * drj) / (mu_edge + EPS); | |
| } |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@cpp/src/umap/densmap.cuh` around lines 347 - 362, The gradient computation
divides by dist_squared, phi_sum entries, and mu_edge which can be zero; to fix,
add small positive guards (e.g., T eps = max(T(1e-12),
std::numeric_limits<T>::epsilon())) and replace direct divides with safe
denominators: use denom_dist = max(dist_squared, eps) for dphi_term, use
denom_phi_sum_k = max(phi_sum[k], eps) and denom_phi_sum_j = max(phi_sum[j],
eps) for q_jk/q_kj, and use denom_mu_edge = max(mu_edge, eps) when dividing the
final value; ensure these guarded variables are used where dist_squared,
phi_sum[k]/[j], and mu_edge appear (symbols: dist_squared, dphi_term, q_jk,
q_kj, phi_sum, mu_edge, drk, drj, outer_scale).
Comment on lines
+113
to
+114
| else if constexpr (metric == DT::LpUnexpanded) | ||
| return pow(acc[0], T(1) / p_val); |
There was a problem hiding this comment.
Guard against p_val == 0 in LpUnexpanded finalization.
If p_val is zero (e.g., misconfigured user input), this line produces Inf. Consider adding a validation check upstream or an epsilon guard here.
🛡️ Suggested guard
else if constexpr (metric == DT::LpUnexpanded)
- return pow(acc[0], T(1) / p_val);
+ return (p_val > T(0)) ? pow(acc[0], T(1) / p_val) : T(0);📝 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
| else if constexpr (metric == DT::LpUnexpanded) | |
| return pow(acc[0], T(1) / p_val); | |
| else if constexpr (metric == DT::LpUnexpanded) | |
| return (p_val > T(0)) ? pow(acc[0], T(1) / p_val) : T(0); |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@cpp/src/umap/edge_metric.cuh` around lines 113 - 114, The LpUnexpanded
finalization currently does pow(acc[0], T(1) / p_val) which will divide by zero
if p_val == 0; update the DT::LpUnexpanded branch in edge_metric.cuh to guard
p_val by clamping it to a small positive epsilon (e.g., T eps =
std::numeric_limits<T>::epsilon() or a tiny constant) before computing
T(1)/p_val, or handle p_val==0 explicitly (return acc[0] or a defined value) so
pow is never called with a division-by-zero; change the expression using the
clamped value (referencing metric, DT::LpUnexpanded, p_val, acc, and T).
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (1)
python/cuml/cuml/manifold/umap/umap.pyx (1)
550-552: Centralize the densMAP defaults.
2.0 / 0.3 / 0.1are now copied across the constructor, C++ param init, and CPU↔GPU serialization. Hoisting them into shared module constants (or a tiny helper) will keep these paths from drifting.Also applies to: 1024-1026, 1059-1062, 1180-1182
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@python/cuml/cuml/manifold/umap/umap.pyx` around lines 550 - 552, Create module-level constants (e.g. DEFAULT_DENS_LAMBDA, DEFAULT_DENS_FRAC, DEFAULT_DENS_VAR_SHIFT) or a tiny helper function and replace the hard-coded literals 2.0/0.3/0.1 wherever dens defaults are set or serialized: use these constants inside the umap class constructor (where dens_lambda/dens_frac/dens_var_shift are read), in the C++ parameter initialization code path, and in the CPU↔GPU serialization/deserialization logic (locations around the occurrences you noted: the constructor, param init, and serialization blocks). Update references to dens_lambda, dens_frac, and dens_var_shift to fall back to the shared defaults via getattr(self, "dens_lambda", DEFAULT_DENS_LAMBDA) (and analogous for the others) so all three places use the same centralized values.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@python/cuml/cuml/manifold/umap/umap.pyx`:
- Around line 885-886: Update the messaging and behavior around densMAP in
umap.pyx so callers of transform() are not silently forwarded to
fit_transform(); instead, when densmap=True and transform() or
inverse_transform() is called on a fitted estimator (or when docs mention
transform availability), raise/emit an explicit error stating densMAP does not
support mapping new data and do not retrain; direct users to the saved training
embedding via the attribute embedding_ for the original embedding and state that
inverse_transform() is unavailable with densMAP. Apply the same change to all
similar messages/branches (including the other occurrences handling densmap in
the file that currently point to fit_transform()).
- Around line 547-565: The densMAP documentation requires metric='euclidean' but
the code and tests allow other metrics; make behavior and docs consistent:
either remove the "metric='euclidean'" restriction from the densMAP docstring,
or enforce it in init_params() by checking the input metric when densmap is True
(raise ValueError if metric != 'euclidean'); if you choose enforcement, perform
the check near where densmap is processed (in init_params() around densmap
handling) and reference the same metric used for KNN (not just the GPU
output-space metric validated in _params_from_cpu()); update params/raises
accordingly.
---
Nitpick comments:
In `@python/cuml/cuml/manifold/umap/umap.pyx`:
- Around line 550-552: Create module-level constants (e.g. DEFAULT_DENS_LAMBDA,
DEFAULT_DENS_FRAC, DEFAULT_DENS_VAR_SHIFT) or a tiny helper function and replace
the hard-coded literals 2.0/0.3/0.1 wherever dens defaults are set or
serialized: use these constants inside the umap class constructor (where
dens_lambda/dens_frac/dens_var_shift are read), in the C++ parameter
initialization code path, and in the CPU↔GPU serialization/deserialization logic
(locations around the occurrences you noted: the constructor, param init, and
serialization blocks). Update references to dens_lambda, dens_frac, and
dens_var_shift to fall back to the shared defaults via getattr(self,
"dens_lambda", DEFAULT_DENS_LAMBDA) (and analogous for the others) so all three
places use the same centralized values.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 830ebf6d-8f15-4ccf-afda-7932eea5a23d
📒 Files selected for processing (1)
python/cuml/cuml/manifold/umap/umap.pyx
Comment on lines
+547
to
+565
| densmap = getattr(self, "densmap", False) | ||
| params.densmap = densmap | ||
| if densmap: | ||
| dens_lambda = getattr(self, "dens_lambda", 2.0) | ||
| dens_frac = getattr(self, "dens_frac", 0.3) | ||
| dens_var_shift = getattr(self, "dens_var_shift", 0.1) | ||
|
|
||
| if dens_lambda < 0: | ||
| raise ValueError(f"Expected `dens_lambda >= 0`, got {dens_lambda}") | ||
| if not (0.0 <= dens_frac <= 1.0): | ||
| raise ValueError(f"Expected `0 <= dens_frac <= 1`, got {dens_frac}") | ||
| if dens_var_shift < 0: | ||
| raise ValueError(f"Expected `dens_var_shift >= 0`, got {dens_var_shift}") | ||
| if is_sparse: | ||
| raise NotImplementedError("densMAP is not supported for sparse inputs") | ||
|
|
||
| params.dens_lambda = dens_lambda | ||
| params.dens_frac = dens_frac | ||
| params.dens_var_shift = dens_var_shift |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect the densMAP metric contract across Python, tests, and CUDA code.
for f in \
python/cuml/cuml/manifold/umap/umap.pyx \
python/cuml/tests/test_umap_densmap.py \
cpp/include/cuml/manifold/umapparams.h \
cpp/src/umap/densmap.cuh \
cpp/src/umap/runner.cuh \
cpp/src/umap/simpl_set_embed/algo.cuh \
cpp/src/umap/simpl_set_embed/optimize_batch_kernel.cuh
do
echo "=== $f ==="
rg -n -C2 'densmap|output_metric|metric' "$f"
doneRepository: rapidsai/cuml
Length of output: 31764
Update the densMAP metric contract to match the implementation.
The documentation at line 879-884 states that densmap=True requires metric='euclidean', but the implementation never enforces this restriction during fitting. Moreover, the test suite (test_umap_densmap.py:119-144) explicitly validates densMAP with multiple metrics (euclidean, l2, sqeuclidean, cosine, manhattan, hellinger, minkowski), indicating the intended behavior is to support non-euclidean metrics.
Remove the metric='euclidean' requirement from the docstring or, if that requirement is actually necessary for numerical correctness, add explicit validation in init_params() at line 547-565 to enforce it. The check in _params_from_cpu() at line 988-993 only validates the GPU output space metric, not the input metric used for KNN, so it does not enforce the documented contract.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@python/cuml/cuml/manifold/umap/umap.pyx` around lines 547 - 565, The densMAP
documentation requires metric='euclidean' but the code and tests allow other
metrics; make behavior and docs consistent: either remove the
"metric='euclidean'" restriction from the densMAP docstring, or enforce it in
init_params() by checking the input metric when densmap is True (raise
ValueError if metric != 'euclidean'); if you choose enforcement, perform the
check near where densmap is processed (in init_params() around densmap handling)
and reference the same metric used for KNN (not just the GPU output-space metric
validated in _params_from_cpu()); update params/raises accordingly.
Comment on lines
+885
to
+886
| ``transform()`` is not supported when ``densmap=True``; use | ||
| ``fit_transform()`` instead. |
There was a problem hiding this comment.
Don't send transform() callers to fit_transform().
On a fitted estimator, fit_transform() retrains instead of mapping into the existing embedding. The docs/error text should say that densMAP does not support transforming new data, point users at embedding_ for the training embedding, and mention that inverse_transform() is also unavailable.
✏️ Suggested wording
- ``transform()`` is not supported when ``densmap=True``; use
- ``fit_transform()`` instead.
+ ``transform()`` is not supported for new data when ``densmap=True``.
+ Use ``embedding_`` to access the fitted training embedding.
+ ``inverse_transform()`` is also not supported when ``densmap=True``.
...
- "transform is not supported for densMAP. "
- "Use fit_transform instead."
+ "transform is not supported for densMAP on new data. "
+ "Use `embedding_` to access the fitted training embedding."Also applies to: 1492-1496, 1636-1637
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@python/cuml/cuml/manifold/umap/umap.pyx` around lines 885 - 886, Update the
messaging and behavior around densMAP in umap.pyx so callers of transform() are
not silently forwarded to fit_transform(); instead, when densmap=True and
transform() or inverse_transform() is called on a fitted estimator (or when docs
mention transform availability), raise/emit an explicit error stating densMAP
does not support mapping new data and do not retrain; direct users to the saved
training embedding via the attribute embedding_ for the original embedding and
state that inverse_transform() is unavailable with densMAP. Apply the same
change to all similar messages/branches (including the other occurrences
handling densmap in the file that currently point to fit_transform()).
Contributor
Author
|
Closing for now as the project has been de-prioritized. |
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
4 participants
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.
Closes #4573.