Project: 018_Learning_to_Discretize_Denoising_Diffusion_ODEs | Function: marginal_lambda (iclr)

system: You are an expert Python programmer implementing algorithms from scientific papers.
Complete the masked function based on the algorithm description.
Follow Python best practices and maintain consistent code style.
Return **only** the complete Python function code.

user: Complete this Python function:

**DESCRIPTION:** Noise schedule lambda function computation implementing log-SNR transformation λt = log(αt) - log(σt) for diffusion ODE parameterization and time step conversion

**KEY TERMS:** {
  "log_mean_coeff": "log(αt) mean coefficient logarithm",
  "log_std": "log(σt) standard deviation logarithm",
  "marginal_log_mean_coeff": "log(αt) computation function",
  "torch.logaddexp": "numerically stable log-sum-exp operation",
  "lambda_t": "λt = log(αt) - log(σt) log-SNR value"
}

**CONTEXT CODE:**             log_alpha_fn = lambda s: torch.log(torch.cos((s + self.cosine_s) / (1. + self.cosine_s) * math.pi / 2.))
            log_alpha_t =  log_alpha_fn(t) - self.cosine_log_alpha_0
            return log_alpha_t
        else:
            raise ValueError("Unsupported noise schedule {}".format(self.schedule))
        
    def dalpha_dt(self, t):
        if self.schedule == 'cosine': # need check!
            return 0.5 * math.pi * torch.sin((t + self.cosine_s) / (1. + self.cosine_s) * math.pi / 2.) / (1. + self.cosine_s)
        else:
            return -0.5 * t * (self.beta_1 - self.beta_0) - 0.5 * self.beta_0
    
    def dsigma_dt(self, t):
        alpha_t = torch.exp(self.marginal_log_mean_coeff(t))
        # we know that sigma_t = sqrt(1 - alpha_t^2)
        # d(sigma_t) / dt = - alpha_t * d(alpha_t) / dt / sqrt(1 - alpha_t^2)
        return - alpha_t * self.dalpha_dt(t) / torch.sqrt(1. - alpha_t ** 2)
    
    def ft(self, t):
        if self.schedule == 'discrete':
            return NoiseScheduleVP.derivative(self.marginal_log_mean_coeff, t)
        dalpha_dt = self.dalpha_dt(t)
        alpha_t = torch.exp(self.marginal_log_mean_coeff(t))
        return dalpha_dt / alpha_t
    
    def gt(self, t):
        if self.schedule == 'discrete':
            marginal_std_at_t = self.marginal_std(t)
            return 2 * marginal_std_at_t * NoiseScheduleVP.derivative(self.marginal_std, t) - 2 * (marginal_std_at_t ** 2 )* self.ft(t)
        
        # gt = sqrt(2sigma_t * dsigma_t / dt - 2 ft(t) * sigma_t^2)
        gt = torch.sqrt(2 * self.marginal_std(t) * self.dsigma_dt(t) - 2 * self.ft(t) * self.marginal_std(t) ** 2)
        return gt 

    def marginal_alpha(self, t): # alpha_t
        """
        Compute alpha_t of a given continuous-time label t in [0, T].
        """
        if not isinstance(t, torch.Tensor):
            t = torch.tensor(t)
        return torch.exp(self.marginal_log_mean_coeff(t))

    def marginal_std(self, t): # sigma_t
        """
        Compute sigma_t of a given continuous-time label t in [0, T].
        """
        if not isinstance(t, torch.Tensor):
            t = torch.tensor(t)
        return torch.sqrt(1. - torch.exp(2. * self.marginal_log_mean_coeff(t)))

    def marginal_lambda(self, t): # lambda_t
        """
        Compute lambda_t = log(alpha_t) - log(sigma_t) of a given continuous-time label t in [0, T].
        """
        raise NotImplementedError("Function implementation has been masked for testing")
    
    def inverse_lambda(self, lamb): # t
        """
        Compute the continuous-time label t in [0, T] of a given half-logSNR lambda_t.
        """
        if not isinstance(lamb, torch.Tensor):
            lamb = torch.tensor(lamb)
            scalar=True
            
        if self.schedule == 'linear':
            tmp = 2. * (self.beta_1 - self.beta_0) * torch.logaddexp(-2. * lamb, torch.zeros((1,)).to(lamb))
            Delta = self.beta_0**2 + tmp
            ret = tmp / (torch.sqrt(Delta) + self.beta_0) / (self.beta_1 - self.beta_0)
        elif self.schedule == 'discrete':
            log_alpha = -0.5 * torch.logaddexp(torch.zeros((1,)).to(lamb.device), -2.0 * lamb)
            t = interpolate_fn(
                log_alpha.reshape((-1, 1)),
                torch.flip(self.log_alpha_array.to(lamb.device), [1]),
                torch.flip(self.t_array.to(lamb.device), [1]),
            )
            ret = t.reshape((-1,))
        elif self.schedule == 'cosine': # cosine
            log_alpha = -0.5 * torch.logaddexp(-2. * lamb, torch.zeros((1,)).to(lamb))
            t_fn = lambda log_alpha_t: torch.arccos(torch.exp(log_alpha_t + self.cosine_log_alpha_0)) * 2. * (1. + self.cosine_s) / math.pi - self.cosine_s
            ret = t_fn(log_alpha)
        else:
            raise ValueError("Unsupported noise schedule {}".format(self.schedule))
        
        return ret
        
    def prior_transformation(self, latents):
        return latents
    
    def numerical_clip_alpha(self, log_alphas, clipped_lambda=-5.1):
        """
        For some beta schedules such as cosine schedule, the log-SNR has numerical isssues.
        We clip the log-SNR near t=T within -5.1 to ensure the stability.
        Such a trick is very useful for diffusion models with the cosine schedule, such as i-DDPM, guided-diffusion and GLIDE.
        """
        log_sigmas = 0.5 * torch.log(1.0 - torch.exp(2.0 * log_alphas))
        lambs = log_alphas - log_sigmas
        idx = torch.searchsorted(torch.flip(lambs, [0]), clipped_lambda)
        if idx > 0:
            log_alphas = log_alphas[:-idx]
        return log_alphas


class NoiseScheduleVE:
    def __init__(
            self,
            schedule='edm', # discrete, heat, edm
            sigma_min=0.002,
            sigma_max=80.,
            N=1000,
        ):
        """Create a wrapper class for the forward SDE (VE type).

        The forward SDE ensures that the condition distribution q_{t|0}(x_t | x_0) = N ( alpha_t * x_0, sigma_t^2 * I ).
        We further define lambda_t = log(alpha_t) - log(sigma_t), which is the half-logSNR (described in the DPM-Solver paper).
        Therefore, we implement the functions for computing alpha_t, sigma_t and lambda_t. For t in [0, T], we have:

            log_alpha_t = self.marginal_log_mean_coeff(t)
            sigma_t = self.marginal_std(t)
            lambda_t = self.marginal_lambda(t)

        Moreover, as lambda(t) is an invertible function, we also support its inverse function:

            t = self.inverse_lambda(lambda_t)

        ===============================================================

        Args:
            sigma_min: A `float` number. The smallest sigma for the VE schedule.
            sigma_max: A `float` number. The largest sigma for the VE schedule.
            N: An `int` number. The number of time steps for the VE schedule.
        Returns:
            A wrapper object of the forward SDE (VE type).
        """

        self.schedule = schedule
        self.sigma_min = sigma_min
        self.sigma_max = sigma_max
        self.N = N
        
        if schedule == 'heat':
            self.T = self.sigma_max ** 2
            self.eps = self.sigma_min ** 2
        elif schedule == 'edm':
            self.T = self.sigma_max
            self.eps = self.sigma_min
        
        self.lambda_max = self.marginal_lambda(self.eps).item()
        self.lambda_min = self.marginal_lambda(self.T).item()

    '''


Complete the function implementation replacing "raise NotImplementedError".