Skip to content

Clamp negative SDE sampler tmp#15843

Open
fallintoplace wants to merge 1 commit into
NVIDIA-NeMo:mainfrom
fallintoplace:fix/sb-sde-sqrt-clamp
Open

Clamp negative SDE sampler tmp#15843
fallintoplace wants to merge 1 commit into
NVIDIA-NeMo:mainfrom
fallintoplace:fix/sb-sde-sqrt-clamp

Conversation

@fallintoplace

Copy link
Copy Markdown

Summary

Clamp the SDE sampler's intermediate tmp term before sqrt() so tiny negative drift cannot turn the update into NaNs.

Root cause

In the SDE branch, the sampler computes tmp = 1 - sigma_t**2 / (sigma_prev**2 + eps) and passes it directly to sqrt(). The shipped schedules stay non-negative in the normal path, but any slight monotonicity drift from finite precision or edge-case overrides can push tmp just below zero and poison the sample with NaNs.

What changed

  • clamp tmp with clamp_min(0) before using it in the SDE update
  • add a focused regression test that constructs adjacent timesteps near the boundary and forces a tiny negative pre-clamp tmp
  • verify the sampler output stays finite and matches the clamped update behavior

Impact

The SDE sampler is more numerically robust around boundary steps and invalid-nearby schedules, without changing the normal monotonic cases.

Validation

  • uv run pytest tests/collections/audio/test_audio_parts_submodules_schroedinger_bridge.py -k "clamps_negative_tmp_before_sqrt" -v
  • uv run pytest tests/collections/audio/test_audio_parts_submodules_schroedinger_bridge.py -v
  • uv run --with black==24.10.0 black --check nemo/collections/audio/parts/submodules/schroedinger_bridge.py tests/collections/audio/test_audio_parts_submodules_schroedinger_bridge.py
  • uv run isort --check-only nemo/collections/audio/parts/submodules/schroedinger_bridge.py tests/collections/audio/test_audio_parts_submodules_schroedinger_bridge.py

@copy-pr-bot

copy-pr-bot Bot commented Jun 27, 2026

Copy link
Copy Markdown

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants