Forward & Backward Diffusion Process

Implement the complete forward and backward diffusion processes used in Denoising Diffusion Probabilistic Models (DDPM).

The forward process gradually adds Gaussian noise to data according to a noise schedule (betas). Given clean data x_0, timestep t, and noise epsilon, compute the noisy sample x_t using the closed-form reparameterization.

The backward (reverse) process denoises the data by predicting and removing noise. Given x_t and a predicted noise from a model, compute x_{t-1} using the posterior distribution parameters.

Implement a function that:

1. Computes x_t from x_0 using the forward diffusion formula
2. Computes x_{t-1} from x_t using the backward diffusion step

The function should handle the special case when t=1 (final denoising step) where no additional noise is added.

Parameters:

• x_0: Original clean data
• betas: Noise schedule array of length T
• timestep: Current timestep t (1-indexed, 1 <= t <= T)
• forward_noise: Random noise for forward process
• predicted_noise: Noise predicted by the denoising model
• backward_noise: Random noise for stochastic backward step (None when t=1)

import numpy as np

def diffusion_process(x_0: np.ndarray, 
                      betas: np.ndarray,
                      timestep: int,
                      forward_noise: np.ndarray,
                      predicted_noise: np.ndarray,
                      backward_noise: np.ndarray = None) -> tuple:
    """
    Implement forward and backward diffusion processes.
    
    Args:
        x_0: Original clean data (any shape)
        betas: Noise schedule array of shape (T,)
        timestep: Current timestep t (1-indexed)
        forward_noise: Noise epsilon for forward diffusion
        predicted_noise: Model's predicted noise for backward diffusion
        backward_noise: Random noise z for stochastic backward step
    
    Returns:
        tuple: (x_t, x_t_minus_1) - noisy and denoised samples
    """
    pass