expose number of diffusion samples, use ds as axis for diffusion samples

chai_lab/chai1.py

@@ -443,6 +443,7 @@ def run_inference(
     # expose some params for easy tweaking
     num_trunk_recycles: int = 3,
     num_diffn_timesteps: int = 200,
+    num_diffn_samples: int = 5,
     seed: int | None = None,
     device: str | None = None,
     low_memory: bool = True,
@@ -470,6 +471,7 @@ def run_inference(
         output_dir=output_dir,
         num_trunk_recycles=num_trunk_recycles,
         num_diffn_timesteps=num_diffn_timesteps,
+        num_diffn_samples=num_diffn_samples,
         seed=seed,
         device=torch_device,
         low_memory=low_memory,
@@ -488,6 +490,8 @@ def run_folding_on_context(
     # expose some params for easy tweaking
     num_trunk_recycles: int = 3,
     num_diffn_timesteps: int = 200,
+    # all diffusion samples come from the same trunk
+    num_diffn_samples: int = 5,
     seed: int | None = None,
     device: torch.device | None = None,
     low_memory: bool,
@@ -679,6 +683,7 @@ def run_folding_on_context(
     )
 
     def _denoise(atom_pos: Tensor, sigma: Tensor, ds: int) -> Tensor:
+        # verified manually that ds dimension can be arbitrary in diff module
         atom_noised_coords = rearrange(
             atom_pos, "(b ds) ... -> b ds ...", ds=ds
         ).contiguous()
@@ -690,7 +695,6 @@ def _denoise(atom_pos: Tensor, sigma: Tensor, ds: int) -> Tensor:
             **static_diffusion_inputs,
         )
 
-    num_diffn_samples = 5  # Fixed at export time
     inference_noise_schedule = InferenceNoiseSchedule(
         s_max=DiffusionConfig.S_tmax,
         s_min=4e-4,