ETS improved tests, tutorial, and Generator"

Signed-off-by: btrentini <brunoxtrentini@gmail.com>

Author: btrentini

sub-packages/bionemo-moco/src/bionemo/moco/schedules/inference_time_schedules.py

@@ -545,9 +545,10 @@ def _hutchinson_divergence(self, t: Tensor, x: Tensor) -> Tensor:
 
         # random vector from the Rademacher distribution
         if self.generator:
-            epsilon = (torch.randint_like(x, 0, 2, generator=self.generator) * 2 - 1).to(x.dtype)
+            epsilon = (torch.randint(0, 2, x.shape, generator=self.generator, device=x.device) * 2 - 1).to(x.dtype)
         else:
             epsilon = (torch.randint_like(x, 0, 2) * 2 - 1).to(x.dtype)
+
 
         v = self.predictor(t, x)
 

sub-packages/bionemo-moco/src/bionemo/moco/schedules/utils.py

@@ -22,3 +22,4 @@ class TimeDirection(Enum):
 
     UNIFIED = "unified"  # Noise(0) --> Data(1)
     DIFFUSION = "diffusion"  # Noise(1) --> Data(0)
+

sub-packages/bionemo-moco/tests/bionemo/moco/schedules/test_inference_schedules.py

@@ -22,6 +22,7 @@
     LinearInferenceSchedule,
     LogInferenceSchedule,
     PowerInferenceSchedule,
+    EntropicInferenceSchedule,
 )
 from bionemo.moco.schedules.utils import TimeDirection
 
@@ -215,46 +216,46 @@ def test_entropic_schedule(timesteps, device, direction):
 
 @pytest.mark.parametrize("device", ["cpu", "cuda"])
 def test_entropic_schedule_reproducibility(device):
-    """Checks that the the EntropicInferenceSchedule produce reproducible results when 
-    a torch.Generator with a fixed seed is provided."""
+    """Checks that the the EntropicInferenceSchedule produce reproducible results when a torch.Generator with a fixed seed is provided."""
     if device == "cuda" and not torch.cuda.is_available():
         pytest.skip("CUDA is not available")
 
     timesteps = 10
     dim = 2
-    predictor = lambda t, x: (2 * t - 1) * x
-    x_0_sampler = lambda bs: torch.randn(bs, dim, device=device)
-    x_1_sampler = lambda bs: torch.randn(bs, dim, device=device)
+    predictor = lambda t, x: t * torch.sin(x)
 
     gen1 = torch.Generator(device=device).manual_seed(42)
+    sampler1 = lambda bs: torch.randn(bs, dim, device=device, generator=gen1)
     scheduler1 = EntropicInferenceSchedule(
         predictor=predictor,
-        x_0_sampler=x_0_sampler,
-        x_1_sampler=x_1_sampler,
+        x_0_sampler=sampler1,
+        x_1_sampler=sampler1,
         nsteps=timesteps,
         device=device,
         generator=gen1,
     )
     schedule1 = scheduler1.generate_schedule()
 
-    # Run again with the same seed...
+    # Run again with the same seed ---
     gen2 = torch.Generator(device=device).manual_seed(42)
+    sampler2 = lambda bs: torch.randn(bs, dim, device=device, generator=gen2)
     scheduler2 = EntropicInferenceSchedule(
         predictor=predictor,
-        x_0_sampler=x_0_sampler,
-        x_1_sampler=x_1_sampler,
+        x_0_sampler=sampler2,
+        x_1_sampler=sampler2,
         nsteps=timesteps,
         device=device,
         generator=gen2,
     )
     schedule2 = scheduler2.generate_schedule()
-
+    
     # Compare again with another seed
     gen3 = torch.Generator(device=device).manual_seed(99)
+    sampler3 = lambda bs: torch.randn(bs, dim, device=device, generator=gen3)
     scheduler3 = EntropicInferenceSchedule(
         predictor=predictor,
-        x_0_sampler=x_0_sampler,
-        x_1_sampler=x_1_sampler,
+        x_0_sampler=sampler3,
+        x_1_sampler=sampler3,
         nsteps=timesteps,
         device=device,
         generator=gen3,
@@ -264,5 +265,5 @@ def test_entropic_schedule_reproducibility(device):
     # Schedules from identical seeds should be identical
     assert torch.allclose(schedule1, schedule2)
 
-    # Schedule from a different seed should be different
+    # Schedules from different seeds should be different
     assert not torch.allclose(schedule1, schedule3)