📝 fix some more docstrings

src/DDPM.jl

@@ -3,7 +3,7 @@ include("Schedulers.jl")
 """
 Denoising Diffusion Probabilistic Models (DDPM) scheduler.
 
-https://arxiv.org/abs/2006.11239
+cf. [[2006.11239] Denoising Diffusion Probabilistic Models](https://arxiv.org/abs/2006).
 """
 struct DDPM{V<:AbstractVector} <: Scheduler
   # number of diffusion steps used to train the model.
@@ -41,14 +41,15 @@ end
 """
 Remove noise from model output using the backward diffusion process.
 
-Args:
+## Input
-  scheduler (`DDPM`): scheduler object.
+  * scheduler (`DDPM`): scheduler to use
-  sample (`AbstractArray`): sample to remove noise from, i.e. model_input.
+  * sample (`AbstractArray`): sample to remove noise from, i.e. model_input
-  model_output (`AbstractArray`): predicted noise from the model.
+  * model_output (`AbstractArray`): predicted noise from the model
-  timesteps (`AbstractArray`): timesteps to remove noise from.
+  * timesteps (`AbstractArray`): timesteps to remove noise from
 
-Returns:
+## Output
-  `AbstractArray`: denoised model output at the given timestep.
+  * pred_prev_sample (`AbstractArray`): denoised sample at t=t-1
+  * x_0_pred (`AbstractArray`): denoised sample at t=0
 """
 function step(
   scheduler::DDPM,

src/Schedulers.jl

@@ -4,13 +4,13 @@ abstract type Scheduler end
 Add noise to clean data using the forward diffusion process.
 
 ## Input
-  * scheduler (`Scheduler`): scheduler to use.
+  * scheduler (`Scheduler`): scheduler to use
-  * clean_data (`AbstractArray`): clean data to add noise to.
+  * clean_data (`AbstractArray`): clean data to add noise to
-  * noise (`AbstractArray`): noise to add to clean data.
+  * noise (`AbstractArray`): noise to add to clean data
-  * timesteps (`AbstractArray`): timesteps used to weight the noise.
+  * timesteps (`AbstractArray`): timesteps used to weight the noise
 
 ## Output
-  * noisy_data (`AbstractArray`): noisy data at the given timesteps.
+  * noisy_data (`AbstractArray`): noisy data at the given timesteps
 """
 function add_noise(
   scheduler::Scheduler,