# \#4 JAX Course - JAX Implementation of DALL·E. - Generate images from a prompt **URL:** \ **Category:** JAX\ **Created:** [December 14, 2021, 5:46pm UTC](https://community.wandb.ai/t/4-jax-course-jax-implementation-of-dall-e-generate-images-from-a-prompt/1532 "2021-12-14T17:46:58Z")\ **Posts on this page:** 14\ **Page:** 1
**Author:** ![ani-safaryan](https://sea2.discourse-cdn.com/flex020/user_avatar/community.wandb.ai/ani-safaryan/32/527_2.png) [@ani-safaryan](https://community.wandb.ai/u/ani-safaryan)\ **Post date:** [December 14, 2021, 5:46pm UTC](https://community.wandb.ai/t/4-jax-course-jax-implementation-of-dall-e-generate-images-from-a-prompt/1532/1 "2021-12-14T17:46:59Z")
🙌🔥We’re thrilled to host Boris Dayma an MLE at W&B He will be teaching us about his work on implementing Dall-e in JAX 👨‍🏫 **How to Join?** 🔗[Register here to join live](https://wandb.zoom.us/webinar/register/WN_L1rCb73XRL-N-j1uHgFQYg) 🎥Recording: [JAX Course - 4. JAX Implementation of DALL·E - YouTube](https://www.youtube.com/watch?v=WKOQ3XaQ620&list=PLD80i8An1OEFy7Yrw8Q8VgG0vpWbd_LLu&index=4) ---
**Author:** ![nkirukauzuegbunam](https://avatars.discourse-cdn.com/v4/letter/n/74df32/32.png) [@nkirukauzuegbunam](https://community.wandb.ai/u/nkirukauzuegbunam)\ **Post date:** [December 18, 2021, 4:18pm UTC](https://community.wandb.ai/t/4-jax-course-jax-implementation-of-dall-e-generate-images-from-a-prompt/1532/2 "2021-12-18T16:18:53Z")
For the inference pipeline: i missed it but why do you generate these many sample candidate prompts? ---
**Author:** ![ramesh](https://sea2.discourse-cdn.com/flex020/user_avatar/community.wandb.ai/ramesh/32/98_2.png) [@ramesh](https://community.wandb.ai/u/ramesh)\ **Post date:** [December 18, 2021, 4:19pm UTC](https://community.wandb.ai/t/4-jax-course-jax-implementation-of-dall-e-generate-images-from-a-prompt/1532/3 "2021-12-18T16:19:31Z")
Training Pipeline seems to be VQGAN Encoder, BART then Cross entropy loss. What part of this is Trained? Was VQGAN Encoder - Trained? Inference Pipeline: VQGAN Decoder seems to be only during inference. Is the Decoder or CLIP Trained? ---
**Author:** ![ramesh](https://sea2.discourse-cdn.com/flex020/user_avatar/community.wandb.ai/ramesh/32/98_2.png) [@ramesh](https://community.wandb.ai/u/ramesh)\ **Post date:** [December 18, 2021, 4:23pm UTC](https://community.wandb.ai/t/4-jax-course-jax-implementation-of-dall-e-generate-images-from-a-prompt/1532/4 "2021-12-18T16:23:31Z")
Training is clear. But how is it connected to Inference? How do you generate multiple candidate images from one BART Decoder Output? Is the ranking / CLIP part learned? ---
**Author:** ![ramesh](https://sea2.discourse-cdn.com/flex020/user_avatar/community.wandb.ai/ramesh/32/98_2.png) [@ramesh](https://community.wandb.ai/u/ramesh)\ **Post date:** [December 18, 2021, 4:25pm UTC](https://community.wandb.ai/t/4-jax-course-jax-implementation-of-dall-e-generate-images-from-a-prompt/1532/5 "2021-12-18T16:25:31Z")
Thanks for the detailed explanation on Generation output. What are possible values for the Outputs - 0-255 in 3 channels? How do you rank the images? Compare the embeddings of input text to generated image? ---
**Author:** ![nkirukauzuegbunam](https://avatars.discourse-cdn.com/v4/letter/n/74df32/32.png) [@nkirukauzuegbunam](https://community.wandb.ai/u/nkirukauzuegbunam)\ **Post date:** [December 18, 2021, 4:26pm UTC](https://community.wandb.ai/t/4-jax-course-jax-implementation-of-dall-e-generate-images-from-a-prompt/1532/6 "2021-12-18T16:26:52Z")
Nice idea to concat a bart encoding of the text, with the fixed vqgan encoding of the image. Does this mean we can extend this to maybe audio like replace vqgan encoding with an audio encoding? or even a graph encoding if that is possible ---
**Author:** ![ramesh](https://sea2.discourse-cdn.com/flex020/user_avatar/community.wandb.ai/ramesh/32/98_2.png) [@ramesh](https://community.wandb.ai/u/ramesh)\ **Post date:** [December 18, 2021, 4:32pm UTC](https://community.wandb.ai/t/4-jax-course-jax-implementation-of-dall-e-generate-images-from-a-prompt/1532/7 "2021-12-18T16:32:55Z")
One general question - Are there things that JAX / Flax enables you to do, that is harder to do in PyTorch? Similarly are there things that are easier in PyTorch that are harder to do in JAX / Flax? ---
**Author:** ![ramesh](https://sea2.discourse-cdn.com/flex020/user_avatar/community.wandb.ai/ramesh/32/98_2.png) [@ramesh](https://community.wandb.ai/u/ramesh)\ **Post date:** [December 18, 2021, 4:37pm UTC](https://community.wandb.ai/t/4-jax-course-jax-implementation-of-dall-e-generate-images-from-a-prompt/1532/8 "2021-12-18T16:37:11Z")
One more related question - Do you use JAX/Flax for any recent experiments or do you prefer to use PyTorch primarily? ---
**Author:** ![ramesh](https://sea2.discourse-cdn.com/flex020/user_avatar/community.wandb.ai/ramesh/32/98_2.png) [@ramesh](https://community.wandb.ai/u/ramesh)\ **Post date:** [December 18, 2021, 4:38pm UTC](https://community.wandb.ai/t/4-jax-course-jax-implementation-of-dall-e-generate-images-from-a-prompt/1532/9 "2021-12-18T16:38:52Z")
Training Question - In CNN / RNN, we might be looking at Gradients or look for Gradient Flow. In Transformers, we primarily only look at Loss value and there are no debugging or issues in Training? What problems you faced in Training DALL-E ---
**Author:** ![nkirukauzuegbunam](https://avatars.discourse-cdn.com/v4/letter/n/74df32/32.png) [@nkirukauzuegbunam](https://community.wandb.ai/u/nkirukauzuegbunam)\ **Post date:** [December 18, 2021, 4:39pm UTC](https://community.wandb.ai/t/4-jax-course-jax-implementation-of-dall-e-generate-images-from-a-prompt/1532/10 "2021-12-18T16:39:42Z")
So, encoded\_images generates some N number of image encodings. What is gen\_top\_k here? Are the encodings rated? ---
**Author:** ![nkirukauzuegbunam](https://avatars.discourse-cdn.com/v4/letter/n/74df32/32.png) [@nkirukauzuegbunam](https://community.wandb.ai/u/nkirukauzuegbunam)\ **Post date:** [December 18, 2021, 4:50pm UTC](https://community.wandb.ai/t/4-jax-course-jax-implementation-of-dall-e-generate-images-from-a-prompt/1532/11 "2021-12-18T16:50:09Z")
So, to init\_weights, can and how do you change the weights to your own pretrained weights for the Bart or the VQGAN so we don’t have to train from scratch ---
**Author:** ![ramesh](https://sea2.discourse-cdn.com/flex020/user_avatar/community.wandb.ai/ramesh/32/98_2.png) [@ramesh](https://community.wandb.ai/u/ramesh)\ **Post date:** [December 18, 2021, 5:06pm UTC](https://community.wandb.ai/t/4-jax-course-jax-implementation-of-dall-e-generate-images-from-a-prompt/1532/12 "2021-12-18T17:06:57Z")
How do you avoid mode collapse and Noise in images? Is this taken care of VQGAN Decoder that it generates reasonable output given a Input vector that you generate from sampled probabilities of BART Decoder? ---
**Author:** ![yuvraj](https://avatars.discourse-cdn.com/v4/letter/y/b5a626/32.png) [@yuvraj](https://community.wandb.ai/u/yuvraj)\ **Post date:** [December 18, 2021, 5:14pm UTC](https://community.wandb.ai/t/4-jax-course-jax-implementation-of-dall-e-generate-images-from-a-prompt/1532/13 "2021-12-18T17:14:04Z")
Sorry, I thought Boris is running on 1 gpu machine, if so how will pmap would help ? Please excuse if I this is a nood question 🙂 ---
**Author:** ![ramesh](https://sea2.discourse-cdn.com/flex020/user_avatar/community.wandb.ai/ramesh/32/98_2.png) [@ramesh](https://community.wandb.ai/u/ramesh)\ **Post date:** [December 20, 2021, 1:57am UTC](https://community.wandb.ai/t/4-jax-course-jax-implementation-of-dall-e-generate-images-from-a-prompt/1532/14 "2021-12-20T01:57:29Z")
From what I understand, pmap will not help on single gpu setting. But it also doesn’t hurt. So if you run the same code in multi-gpu setup, it will benefit without changing any part of the code.