[Flax] 🚨 0.7.0 not working 🚨

[camenduru]

Describe the bug

Reproduction

No response

Logs

No response

System Info

TPU v3-8

[camenduru]

I am still testing...

[camenduru]

0.7.0 🦆 💀 😢

[camenduru]

Maybe related 17 hours ago Jaxlib release v0.3.24 also this new version using more memory 😭

[patrickvonplaten]

Thanks for notifying us here so quickly! I'm not on my computer at the moment - but will try to fix it by tonight!

Could you double check at which commit of main results start to become incorrect?
I suspect that we messed up something between 0.6.0 and 0.7.0 here - also @pcuenca @patil-suraj here

[camenduru]

accelerate sus

[pcuenca]

Thanks @camenduru! I'm not at my computer either, but will look into it later today.

[camenduru]

🚓🚔🚨 0b61cea 🚨🚔🚓

[pcuenca]

Thanks, @camenduru, fix incoming!

[pcuenca]

@camenduru We have a fix candidate. We'll probably merge it soon, but meanwhile if you want to give it a go, you can install diffusers from the PR branch like this:

pip install https://github.com/huggingface/diffusers.git@fix-flax-time-embeddings

If you do test it, we'd appreciate if you'd let us know how it goes :)