Fix SBNoiseSchedule t_max accessors

nemo/collections/audio/parts/submodules/schroedinger_bridge.py

@@ -44,6 +44,7 @@ def __init__(
         eps: float = 1e-8,
     ):
         super().__init__()
+        self.register_buffer('_time_reference', torch.empty(0), persistent=False)
 
         # min and max time
         if time_min < 0:
@@ -86,17 +87,30 @@ def generate_time(self, size: int, device: torch.device) -> torch.Tensor:
         time = torch.rand(size, device=device) * self.time_delta + self.time_min
         return time
 
+    def _time_max_like(self, ref: Optional[torch.Tensor] = None) -> torch.Tensor:
+        """Return `time_max` on the same device and dtype as `ref`."""
+        if ref is None:
+            ref = self._time_reference
+
+        return ref.new_tensor([self.time_max])
+
     @property
     def alpha_t_max(self):
         """Return alpha_t at t_max."""
-        t_max = torch.tensor([self.time_max], device=alpha.device)
-        return self.alpha(t_max)
+        return self.alpha(self._time_max_like())
+
+    def alpha_t_max_like(self, ref: torch.Tensor) -> torch.Tensor:
+        """Return alpha_t at t_max on the same device and dtype as `ref`."""
+        return self.alpha(self._time_max_like(ref))
 
     @property
     def sigma_t_max(self):
         """Return sigma_t at t_max."""
-        t_max = torch.tensor([self.time_max], device=alpha.device)
-        return self.sigma(t_max)
+        return self.sigma(self._time_max_like())
+
+    def sigma_t_max_like(self, ref: torch.Tensor) -> torch.Tensor:
+        """Return sigma_t at t_max on the same device and dtype as `ref`."""
+        return self.sigma(self._time_max_like(ref))
 
     @abstractmethod
     def f(self, time: torch.Tensor) -> torch.Tensor:
@@ -147,7 +161,7 @@ def alpha_bar_from_alpha(self, alpha: torch.Tensor) -> (torch.Tensor, torch.Tens
         Returns:
             Tensors the same size as alpha, representing alpha_bar and alpha_t_max.
         """
-        alpha_t_max = self.alpha(torch.tensor([self.time_max], device=alpha.device))
+        alpha_t_max = self.alpha_t_max_like(alpha)
         alpha_bar = alpha / (alpha_t_max + self.eps)
         return alpha_bar, alpha_t_max
 
@@ -189,7 +203,7 @@ def sigma_bar_from_sigma(self, sigma: torch.Tensor) -> (torch.Tensor, torch.Tens
         Returns:
             Tensors the same size as sigma, representing sigma_bar and sigma_t_max.
         """
-        sigma_t_max = self.sigma(torch.tensor([self.time_max], device=sigma.device))
+        sigma_t_max = self.sigma_t_max_like(sigma)
         sigma_bar_sq = sigma_t_max**2 - sigma**2
         return torch.sqrt(sigma_bar_sq + self.eps), sigma_t_max
 

tests/collections/audio/test_audio_parts_submodules_schroedinger_bridge.py

@@ -19,6 +19,40 @@
 from nemo.collections.audio.parts.submodules.schroedinger_bridge import SBNoiseScheduleVE, SBNoiseScheduleVP, SBSampler
 
 NUM_STEPS = [1, 5, 10, 20, 100]
+DEVICES = [
+    pytest.param(torch.device("cpu"), id="cpu"),
+    pytest.param(
+        torch.device("cuda"),
+        marks=pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available"),
+        id="cuda",
+    ),
+]
+
+
+@pytest.mark.unit
+@pytest.mark.parametrize("dtype", [torch.float32, torch.float64], ids=["float32", "float64"])
+@pytest.mark.parametrize("device", DEVICES)
+def test_sb_noise_schedule_t_max_accessors_match_reference_tensor(device, dtype):
+    noise_schedule = SBNoiseScheduleVE(k=2.0, c=0.5, num_steps=5).to(device=device, dtype=dtype)
+    ref = torch.ones(1, device=device, dtype=dtype)
+
+    alpha_t_max = noise_schedule.alpha_t_max
+    sigma_t_max = noise_schedule.sigma_t_max
+    alpha_t_max_like = noise_schedule.alpha_t_max_like(ref)
+    sigma_t_max_like = noise_schedule.sigma_t_max_like(ref)
+
+    expected_time_max = ref.new_tensor([noise_schedule.time_max])
+    expected_alpha_t_max = noise_schedule.alpha(expected_time_max)
+    expected_sigma_t_max = noise_schedule.sigma(expected_time_max)
+
+    for value in (alpha_t_max, sigma_t_max, alpha_t_max_like, sigma_t_max_like):
+        assert value.device == device
+        assert value.dtype == dtype
+
+    torch.testing.assert_close(alpha_t_max, expected_alpha_t_max)
+    torch.testing.assert_close(sigma_t_max, expected_sigma_t_max)
+    torch.testing.assert_close(alpha_t_max_like, expected_alpha_t_max)
+    torch.testing.assert_close(sigma_t_max_like, expected_sigma_t_max)
 
 
 @pytest.mark.parametrize("num_steps", NUM_STEPS)