fix: address various common pipeline bugs

- models/resnet_flax: add missing bias_init partition bounds to Conv layers
- max_utils: use dict.get() for flash_block_sizes to prevent KeyErrors
- maxdiffusion_utils: simplify VAE PyArrow latency encoding loop to prevent pad dropping and out-of-bounds loop references
- models/attention_flax: allow custom flash_block_sizes to override cross-attention defaults

Author: Perseus14

src/maxdiffusion/max_utils.py

@@ -520,20 +520,20 @@ def get_flash_block_sizes(config):
       max_logging.log(
           "Tokamax kernel specified, Note: Tokamax only supports fused backward kernel."
           "Hence following flash block properties specified will be ignored:"
-          f"block_q: {user_block_sizes['block_q']},"
+          f"block_q: {user_block_sizes.get('block_q')},"
           f"block_q_dq: {user_block_sizes.get('block_q_dq')},"
           f"block_kv_dq: {user_block_sizes.get('block_kv_dq')},"
           f"use_fused_bwd_kernel: {user_block_sizes.get('use_fused_bwd_kernel')}"
       )
     flash_block_sizes = splash_attention_kernel.BlockSizes(
-        block_q=user_block_sizes.get("block_q_dkv", user_block_sizes["block_kv"])
+        block_q=user_block_sizes.get("block_q_dkv", user_block_sizes.get("block_kv"))
         if attention_is_tokamax
-        else user_block_sizes["block_q"],
-        block_kv_compute=user_block_sizes["block_kv_compute"],
-        block_kv=user_block_sizes["block_kv"],
-        block_q_dkv=user_block_sizes["block_q_dkv"],
-        block_kv_dkv=user_block_sizes["block_kv_dkv"],
-        block_kv_dkv_compute=user_block_sizes["block_kv_dkv_compute"],
+        else user_block_sizes.get("block_q"),
+        block_kv_compute=user_block_sizes.get("block_kv_compute", user_block_sizes.get("block_kv")),
+        block_kv=user_block_sizes.get("block_kv"),
+        block_q_dkv=user_block_sizes.get("block_q_dkv", user_block_sizes.get("block_q")),
+        block_kv_dkv=user_block_sizes.get("block_kv_dkv", user_block_sizes.get("block_kv")),
+        block_kv_dkv_compute=user_block_sizes.get("block_kv_dkv_compute", user_block_sizes.get("block_kv")),
         block_q_dq=None if attention_is_tokamax else value_or_none(user_block_sizes, "block_q_dq"),
         block_kv_dq=None if attention_is_tokamax else value_or_none(user_block_sizes, "block_kv_dq"),
         use_fused_bwd_kernel=True if attention_is_tokamax else value_or_none(user_block_sizes, "use_fused_bwd_kernel"),

src/maxdiffusion/maxdiffusion_utils.py

@@ -114,10 +114,15 @@ def transform_images(
   if p_vae_apply:
     tensor_list = np.stack(tensor_list)
     ds_length = tensor_list.shape[0]
-    iters = ds_length // global_batch_size
-    latents_list = []
     local_batch_size = global_batch_size // jax.device_count()
-    for i in range(0, iters * global_batch_size, local_batch_size):
+
+    pad_len = (local_batch_size - (ds_length % local_batch_size)) % local_batch_size
+    if pad_len > 0:
+      pad_tensor = np.zeros((pad_len,) + tensor_list.shape[1:], dtype=tensor_list.dtype)
+      tensor_list = np.concatenate([tensor_list, pad_tensor], axis=0)
+
+    latents_list = []
+    for i in range(0, tensor_list.shape[0], local_batch_size):
       sample_rng, rng = jax.random.split(rng)
       latents = p_vae_apply(tensor_list[i : i + local_batch_size], sample_rng)
       latents_list.append(latents)
@@ -126,14 +131,7 @@ def transform_images(
     b1, b2, c, l1, l2 = latents_list.shape
     latents_list = np.reshape(latents_list, (b1 * b2, c, l1, l2))
 
-    # TODO (Juan Acevedo): do last iteration, its required for the Pyarrow dataset
-    # to not break due to items being fewer than expected. Is there a better way?
-    if tensor_list[i + local_batch_size :].shape[0] != 0:
-      sample_rng, rng = jax.random.split(rng)
-      latents = p_vae_apply(tensor_list[i + local_batch_size :], sample_rng)
-      examples[pixel_ids_key] = np.append(latents_list, latents, axis=0)
-    else:
-      examples[pixel_ids_key] = latents_list
+    examples[pixel_ids_key] = latents_list[:ds_length]
   else:
     examples[pixel_ids_key] = tf.stack(tensor_list)
 

src/maxdiffusion/models/attention_flax.py

@@ -250,19 +250,17 @@ def _tpu_flash_attention(
     kv_max_block_size = ((key.shape[1] + 127) // 128) * 128
   else:
     kv_max_block_size = q_max_block_size
-  # ensure that for cross attention we override the block sizes.
-  if flash_block_sizes and key.shape[1] == query.shape[1]:
+  if flash_block_sizes:
     block_sizes = flash_block_sizes
   else:
-    block_size_q = flash_block_sizes.block_q if flash_block_sizes else q_max_block_size
     block_sizes = splash_attention_kernel.BlockSizes(
-        block_q=block_size_q,
+        block_q=q_max_block_size,
         block_kv_compute=min(kv_max_block_size, key.shape[2]),
         block_kv=min(kv_max_block_size, key.shape[2]),
-        block_q_dkv=block_size_q,
+        block_q_dkv=q_max_block_size,
         block_kv_dkv=min(kv_max_block_size, key.shape[2]),
         block_kv_dkv_compute=min(kv_max_block_size, query.shape[2]),
-        block_q_dq=None if attention_kernel == "tokamax_flash" else block_size_q,
+        block_q_dq=None if attention_kernel == "tokamax_flash" else q_max_block_size,
         block_kv_dq=None if attention_kernel == "tokamax_flash" else min(kv_max_block_size, query.shape[2]),
         use_fused_bwd_kernel=True if attention_kernel == "tokamax_flash" else False,
     )

src/maxdiffusion/models/resnet_flax.py

@@ -51,6 +51,7 @@ def setup(self):
         kernel_init=nn.with_logical_partitioning(
             nn.initializers.lecun_normal(), ("keep_1", "keep_2", "conv_in", "conv_out")
         ),
+        bias_init=nn.with_logical_partitioning(nn.initializers.zeros, ("conv_out",)),
         precision=self.precision,
     )
 
@@ -85,6 +86,7 @@ def setup(self):
         kernel_init=nn.with_logical_partitioning(
             nn.initializers.lecun_normal(), ("keep_1", "keep_2", "conv_in", "conv_out")
         ),
+        bias_init=nn.with_logical_partitioning(nn.initializers.zeros, ("conv_out",)),
         precision=self.precision,
     )
 
@@ -131,6 +133,7 @@ def setup(self):
           kernel_init=nn.with_logical_partitioning(
               nn.initializers.lecun_normal(), ("keep_1", "keep_2", "conv_in", "conv_out")
           ),
+          bias_init=nn.with_logical_partitioning(nn.initializers.zeros, ("conv_out",)),
           precision=self.precision,
       )
     out_channels = self.in_channels if self.out_channels is None else self.out_channels
@@ -144,6 +147,7 @@ def setup(self):
         kernel_init=nn.with_logical_partitioning(
             nn.initializers.lecun_normal(), ("keep_1", "keep_2", "conv_in", "conv_out")
         ),
+        bias_init=nn.with_logical_partitioning(nn.initializers.zeros, ("conv_out",)),
         precision=self.precision,
     )
 
@@ -159,6 +163,7 @@ def setup(self):
         kernel_init=nn.with_logical_partitioning(
             nn.initializers.lecun_normal(), ("keep_1", "keep_2", "conv_in", "conv_out")
         ),
+        bias_init=nn.with_logical_partitioning(nn.initializers.zeros, ("conv_out",)),
         precision=self.precision,
     )