added pad_time function

sub-packages/bionemo-moco/documentation.md

@@ -1621,6 +1621,30 @@ Generate the time schedule as a tensor.
 - `nsteps` _Optioanl[int]_ - Number of time steps. If None, uses the value from initialization.
 - `device` _Optional[str]_ - Device to place the schedule on (default is "cpu").
 
+<a id="mocoschedulesinference_time_schedulesInferenceSchedulepad_time"></a>
+
+#### pad\_time
+
+```python
+def pad_time(n_samples: int,
+             scalar_time: Float,
+             device: Optional[Union[str, torch.device]] = None) -> Tensor
+```
+
+Creates a tensor of shape (n_samples,) filled with a scalar time value.
+
+**Arguments**:
+
+  - n_samples (int): The desired dimension of the output tensor.
+  - scalar_time (Float): The scalar time value to fill the tensor with.
+  - device (Optional[Union[str, torch.device]], optional):
+  The device to place the tensor on. Defaults to None, which uses the default device.
+
+
+**Returns**:
+
+  - Tensor: A tensor of shape (n_samples,) filled with the scalar time value.
+
 <a id="mocoschedulesinference_time_schedulesContinuousInferenceSchedule"></a>
 
 ## ContinuousInferenceSchedule Objects
@@ -2906,7 +2930,7 @@ A Continuous Flow Matching interpolant.
 >>> from bionemo.bionemo.moco.distributions.prior.continuous.gaussian import GaussianPrior
 >>> from bionemo.bionemo.moco.distributions.time.uniform import UniformTimeDistribution
 >>> from bionemo.bionemo.moco.interpolants.continuous_time.continuous.continuous_flow_matching import ContinuousFlowMatcher
->>> from bionemo.bionemo.moco.schedules.inference_time_schedules import LinearTimeSchedule
+>>> from bionemo.bionemo.moco.schedules.inference_time_schedules import LinearInferenceSchedule
 
 flow_matcher = ContinuousFlowMatcher(
     time_distribution = UniformTimeDistribution(...),
@@ -2929,8 +2953,9 @@ for epoch in range(1000):
 
 # Generation
 x_pred = flow_matcher.sample_prior(data.shape)
-for t in LinearTimeSchedule(...).generate_schedule():
-    time = torch.full((batch_size,), t)
+inference_sched = LinearInferenceSchedule(...)
+for t in inference_sched.generate_schedule():
+    time = inference_sched.pad_time(x_pred.shape[0], t)
     u_hat = model(x_pred, time)
     x_pred = flow_matcher.step(u_hat, x_pred, time)
 return x_pred

sub-packages/bionemo-moco/examples/discrete_data_interpolant_tutorial.ipynb

@@ -231,7 +231,7 @@
    "outputs": [],
    "source": [
     "for dt, t in zip(dts, ts):\n",
-    "    t = torch.full((xt.shape[0],), t).to(DEVICE)\n",
+    "    t = schedule.pad_time(num_samples, t, DEVICE)\n",
     "    logits = model(xt, t)\n",
     "    xt = dfm.step(logits, t, xt, dt, stochasticity=0)"
    ]

sub-packages/bionemo-moco/src/bionemo/moco/interpolants/continuous_time/continuous/continuous_flow_matching.py

@@ -55,7 +55,7 @@ class ContinuousFlowMatcher(Interpolant):
     >>> from bionemo.moco.distributions.prior.continuous.gaussian import GaussianPrior
     >>> from bionemo.moco.distributions.time.uniform import UniformTimeDistribution
     >>> from bionemo.moco.interpolants.continuous_time.continuous.continuous_flow_matching import ContinuousFlowMatcher
-    >>> from bionemo.moco.schedules.inference_time_schedules import LinearTimeSchedule
+    >>> from bionemo.moco.schedules.inference_time_schedules import LinearInferenceSchedule
 
     flow_matcher = ContinuousFlowMatcher(
         time_distribution = UniformTimeDistribution(...),
@@ -78,8 +78,9 @@ class ContinuousFlowMatcher(Interpolant):
 
     # Generation
     x_pred = flow_matcher.sample_prior(data.shape)
-    for t in LinearTimeSchedule(...).generate_schedule():
-        time = torch.full((batch_size,), t)
+    inference_sched = LinearInferenceSchedule(...)
+    for t in inference_sched.generate_schedule():
+        time = inference_sched.pad_time(x_pred.shape[0], t)
         u_hat = model(x_pred, time)
         x_pred = flow_matcher.step(u_hat, x_pred, time)
     return x_pred

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

@@ -67,6 +67,22 @@ def generate_schedule(
         """
         pass
 
+    def pad_time(
+        self, n_samples: int, scalar_time: Float, device: Optional[Union[str, torch.device]] = None
+    ) -> Tensor:
+        """Creates a tensor of shape (n_samples,) filled with a scalar time value.
+
+        Args:
+            n_samples (int): The desired dimension of the output tensor.
+            scalar_time (Float): The scalar time value to fill the tensor with.
+            device (Optional[Union[str, torch.device]], optional):
+                The device to place the tensor on. Defaults to None, which uses the default device.
+
+        Returns:
+            Tensor: A tensor of shape (n_samples,) filled with the scalar time value.
+        """
+        return torch.full((n_samples,), fill_value=scalar_time).to(device)
+
 
 class ContinuousInferenceSchedule(InferenceSchedule):
     """A base class for continuous time inference schedules."""