make_image_grid(images, 2, 4)   How does parallelization work? The Flax pipeline in 🤗 Diffusers automatically compiles the model and runs it in parallel on all available devices. Let’s take a closer look at how that process works. JAX parallelization can be done in multiple ways. The easiest one revolves around using the jax.pmap function to achieve single-program multiple-data (SPMD) parallelization. It means running several copies of the same code, each on different data inputs. More sophisticated approaches are possible, and you can go over to the JAX documentation to explore this topic in more detail if you are interested! jax.pmap does two things: Compiles (or ”jits”) the code which is similar to jax.jit(). This does not happen when you call pmap, and only the first time the pmapped function is called. Ensures the compiled code runs in parallel on all available devices. To demonstrate, call pmap on the pipeline’s _generate method (this is a private method that generates images and may be renamed or removed in future releases of 🤗 Diffusers):   Copied p_generate = pmap(pipeline._generate) After calling pmap, the prepared function p_generate will: Make a copy of the underlying function, pipeline._generate, on each device. Send each device a different portion of the input arguments (this is why it’s necessary to call the shard function). In this case, prompt_ids has shape (8, 1, 77, 768) so the array is split into 8 and each copy of _generate receives an input with shape (1, 77, 768). The most important thing to pay attention to here is the batch size (1 in this example), and the input dimensions that make sense for your code. You don’t have to change anything else to make the code work in parallel. The first time you call the pipeline takes more time, but the calls afterward are much faster. The block_until_ready function is used to correctly measure inference time because JAX uses asynchronous dispatch and returns control to the Python loop as soon as it can. You don’t need to use that in your code; blocking occurs automatically when you want to use the result of a computation that has not yet been materialized.   Copied %%time