[Flax] fix Flax scheduler#564

kashif merged 6 commits into

huggingface:mainhuggingface/diffusers:mainfrom

kashif:flax-scheduler-fixkashif/diffusers:flax-scheduler-fixCopy head branch name to clipboard