Flax documentation #589

younesbelkada · 2022-09-20T17:15:03Z

What does this PR do?

This PR adds more documentation on Flax modules
I am planning to add documentation only on modules that are documented on PyTorch modeling files, this includes

Inside attention_flax.py file:

Inside embeddings_flax.py file:

FlaxTimestepEmbedding
FlaxTimesteps

Inside unet_2d_condition_flax.py file:

Inside vae_flax.py file:

Questions

Do we put the paragraph

This model is a Flax Linen [flax.linen.Module](https://flax.readthedocs.io/en/latest/flax.linen.html#module)
    subclass. Use it as a regular Flax linen Module and refer to the Flax documentation for all matter related to
    general usage and behavior.

    Finally, this model supports inherent JAX features such as:
    - [Just-In-Time (JIT) compilation](https://jax.readthedocs.io/en/latest/jax.html#just-in-time-compilation-jit)
    - [Automatic Differentiation](https://jax.readthedocs.io/en/latest/jax.html#automatic-differentiation)
    - [Vectorization](https://jax.readthedocs.io/en/latest/jax.html#vectorization-vmap)
    - [Parallelization](https://jax.readthedocs.io/en/latest/jax.html#parallelization-pmap)

on each nn.Module as we do it in transformers or it's too redundant?

is https://arxiv.org/abs/2103.06104 the correct reference to transformer Unet?

cc @patrickvonplaten @mishig25

HuggingFaceDocBuilderDev · 2022-09-20T17:18:41Z

The documentation is not available anymore as the PR was closed or merged.

mishig25 · 2022-09-20T18:36:57Z

I put them inside the docs page ea8cbdb

should be clear what's going on here

for example, you can see here now that flax docs are rendered in

mishig25 · 2022-09-20T18:56:28Z

to preview locally (quikcer):

Install doc-builder on your system if not already:

git clone https://github.com/huggingface/doc-builder.git
cd doc-builder
pip install -e .

then run:

doc-builder preview diffusers ~/Desktop/diffusers/docs/source/

(obv, change the path based on your system)

younesbelkada · 2022-09-20T20:42:30Z

src/diffusers/models/vae_flax.py

+            Tuple containing the number of output channels for each block
+        layers_per_block (:obj:`int`, *optional*, defaults to `2`):
+            Number of Resnet layer for each block
+        norm_num_groups (:obj:`int`, *optional*, defaults to `2`):


This argument seems to be never used 🤔 Seems to be corresponding to resnet_groups in vae.py but not sure about here

I opened #621 to fix this separately.

Thanks a lot @pcuenca 💪

younesbelkada · 2022-09-20T20:42:48Z

src/diffusers/models/vae_flax.py

+            Activation function
+        latent_channels (:obj:`int`, *optional*, defaults to `4`):
+            Latent space channels
+        norm_num_groups (:obj:`int`, *optional*, defaults to `2`):


Same comment as above

src/diffusers/models/attention_flax.py

src/diffusers/models/vae_flax.py

src/diffusers/models/attention_flax.py

src/diffusers/models/vae_flax.py

younesbelkada · 2022-09-21T10:41:39Z

cc @patil-suraj with #595 being merged I think this is ready for review 🔥 (I cannot somehow request review here)

patrickvonplaten · 2022-09-22T16:39:58Z

Wow super cool!

patrickvonplaten · 2022-09-22T16:41:01Z

Wow this is super cool!

patrickvonplaten

Great effort @younesbelkada ! Looks good to me for merge

pcuenca

Amazing work! Very comprehensive and detailed, thanks a lot!

pcuenca · 2022-09-23T10:05:05Z

src/diffusers/models/unet_2d_condition_flax.py

+    - [Just-In-Time (JIT) compilation](https://jax.readthedocs.io/en/latest/jax.html#just-in-time-compilation-jit)
+    - [Automatic Differentiation](https://jax.readthedocs.io/en/latest/jax.html#automatic-differentiation)
+    - [Vectorization](https://jax.readthedocs.io/en/latest/jax.html#vectorization-vmap)
+    - [Parallelization](https://jax.readthedocs.io/en/latest/jax.html#parallelization-pmap)


👍 ❤️

src/diffusers/models/vae_flax.py

Co-authored-by: Pedro Cuenca <[email protected]>

younesbelkada · 2022-09-23T10:43:44Z

Thank you very much @patrickvonplaten @pcuenca ! 💯
Let me know if you see anything else that we can add, maybe more references to academic papers for some architectures?

pcuenca · 2022-09-23T10:47:52Z

I think it already looks great :) But if you think it's useful to add more references just go for it and merge when you are done :)

younesbelkada · 2022-09-23T10:56:42Z

Perfect looks good to me too! I can't somehow merge nor approve, I think that I don't have the rights on this repo so feel free to merge whenever you want 💪 🔥

* documenting `attention_flax.py` file * documenting `embeddings_flax.py` * documenting `unet_blocks_flax.py` * Add new objs to doc page * document `vae_flax.py` * Apply suggestions from code review * modify `unet_2d_condition_flax.py` * make style * Apply suggestions from code review * make style * Apply suggestions from code review * fix indent * fix typo * fix indent unet * Update src/diffusers/models/vae_flax.py * Apply suggestions from code review Co-authored-by: Pedro Cuenca <[email protected]> Co-authored-by: Mishig Davaadorj <[email protected]> Co-authored-by: Pedro Cuenca <[email protected]>

documenting attention_flax.py file

3f8a55c

younesbelkada and others added 3 commits September 20, 2022 17:24

documenting embeddings_flax.py

f4ecb28

documenting unet_blocks_flax.py

a840591

Add new objs to doc page

ea8cbdb

younesbelkada and others added 2 commits September 20, 2022 21:48

Merge branch 'main' into flax-documentation

53bb433

document vae_flax.py

1341223

younesbelkada commented Sep 20, 2022

View reviewed changes

younesbelkada and others added 5 commits September 20, 2022 22:53

Apply suggestions from code review

7b97067

modify unet_2d_condition_flax.py

edc00f0

make style

dae3064

Apply suggestions from code review

89cd2f3

make style

ce35697

younesbelkada commented Sep 20, 2022

View reviewed changes

src/diffusers/models/vae_flax.py Outdated Show resolved Hide resolved

src/diffusers/models/vae_flax.py Outdated Show resolved Hide resolved

younesbelkada commented Sep 20, 2022

View reviewed changes

src/diffusers/models/attention_flax.py Outdated Show resolved Hide resolved

src/diffusers/models/attention_flax.py Outdated Show resolved Hide resolved

younesbelkada and others added 4 commits September 20, 2022 23:24

Apply suggestions from code review

04ce16e

fix indent

8c52c02

fix typo

29d643e

fix indent unet

309f445

younesbelkada force-pushed the flax-documentation branch from 35716dd to 309f445 Compare September 20, 2022 21:38

younesbelkada commented Sep 21, 2022

View reviewed changes

src/diffusers/models/vae_flax.py Outdated Show resolved Hide resolved

younesbelkada added 2 commits September 21, 2022 12:40

Update src/diffusers/models/vae_flax.py

fb5ee88

Merge branch 'main' into flax-documentation

5880bd5

patrickvonplaten approved these changes Sep 22, 2022

View reviewed changes

pcuenca approved these changes Sep 23, 2022

View reviewed changes

pcuenca mentioned this pull request Sep 23, 2022

Flax: trickle down norm_num_groups #621

Closed

Apply suggestions from code review

fc52b2c

Co-authored-by: Pedro Cuenca <[email protected]>

pcuenca merged commit 8b0be93 into huggingface:main Sep 23, 2022

Flax documentation #589

Flax documentation #589

Uh oh!

Conversation

younesbelkada commented Sep 20, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

Questions

Uh oh!

HuggingFaceDocBuilderDev commented Sep 20, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

mishig25 commented Sep 20, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

mishig25 commented Sep 20, 2022

Uh oh!

younesbelkada Sep 20, 2022

Choose a reason for hiding this comment

Uh oh!

pcuenca Sep 23, 2022

Choose a reason for hiding this comment

Uh oh!

younesbelkada Sep 23, 2022

Choose a reason for hiding this comment

Uh oh!

younesbelkada Sep 20, 2022

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

younesbelkada commented Sep 21, 2022

Uh oh!

patrickvonplaten commented Sep 22, 2022

Uh oh!

patrickvonplaten commented Sep 22, 2022

Uh oh!

patrickvonplaten left a comment

Choose a reason for hiding this comment

Uh oh!

pcuenca left a comment

Choose a reason for hiding this comment

Uh oh!

pcuenca Sep 23, 2022

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

younesbelkada commented Sep 23, 2022

Uh oh!

pcuenca commented Sep 23, 2022

Uh oh!

younesbelkada commented Sep 23, 2022

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

younesbelkada commented Sep 20, 2022 •

edited

Loading

HuggingFaceDocBuilderDev commented Sep 20, 2022 •

edited

Loading

mishig25 commented Sep 20, 2022 •

edited

Loading