class FlaxUNet2DConditionOutput(BaseOutput):
    """
    The output of [`FlaxUNet2DConditionModel`].

    Args:
        sample (`jnp.ndarray` of shape `(batch_size, num_channels, height, width)`):
            The hidden states output conditioned on `encoder_hidden_states` input. Output of last layer of model.
    """

    sample: jnp.ndarray
/Users/nielsrogge/Documents/python_projecten/diffusers/src/diffusers/models/unets/unet_2d_condition_flax.py