Merge branch 'main' into flux_impl

Author: jfacevedo-google

src/maxdiffusion/generate_sdxl.py

@@ -305,7 +305,7 @@ def run(config):
     _ = [stack.enter_context(nn.intercept_methods(interceptor)) for interceptor in lora_interceptors]
     images = p_run_inference(states).block_until_ready()
   print("inference time: ", (time.time() - s))
-  images = jax.experimental.multihost_utils.process_allgather(images)
+  images = jax.experimental.multihost_utils.process_allgather(images, tiled=True)
   numpy_images = np.array(images)
   images = VaeImageProcessor.numpy_to_pil(numpy_images)
   for i, image in enumerate(images):

tests/schedulers/test_scheduler_flax.py

@@ -335,8 +335,8 @@ def test_full_loop_no_noise(self):
         result_mean = jnp.mean(jnp.abs(sample))
 
         if jax_device == "tpu":
-            assert abs(result_sum - 251.26245) < 1e-2
-            assert abs(result_mean - 0.32716465) < 1e-3
+            assert abs(result_sum - 257.2727) < 1e-2
+            assert abs(result_mean - 0.3349905) < 1e-3
         else:
             assert abs(result_sum - 255.1113) < 1e-2
             assert abs(result_mean - 0.332176) < 1e-3