ional, defaults to True) —
Whether to flip the sin to cos in the time embedding.


freq_shift (int, optional, defaults to 0) — The frequency shift to apply to the time embedding.


use_memory_efficient_attention (bool, optional, defaults to False) —
enable memory efficient attention https://arxiv.org/abs/2112.05682



FlaxUNet2DConditionModel is a conditional 2D UNet model that takes in a noisy sample, conditional state, and a
timestep and retu