class FlaxSdeVeOutput(FlaxSchedulerOutput):
    """
    Output class for the ScoreSdeVeScheduler's step function output.

    Args:
        state (`ScoreSdeVeSchedulerState`):
        prev_sample (`jnp.ndarray` of shape `(batch_size, num_channels, height, width)` for images):
            Computed sample (x_{t-1}) of previous timestep. `prev_sample` should be used as next model input in the
            denoising loop.
        prev_sample_mean (`jnp.ndarray` of shape `(batch_size, num_channels, height, width)` for images):
            Mean averaged `prev_sample`. Same as `prev_sample`, only mean-averaged over previous timesteps.
    """

    state: ScoreSdeVeSchedulerState
    prev_sample: jnp.ndarray
    prev_sample_mean: Optional[jnp.ndarray] = None
/Users/nielsrogge/Documents/python_projecten/diffusers/src/diffusers/schedulers/scheduling_sde_ve_flax.py