From e92264486235e3795ff320588d4f818661c73888 Mon Sep 17 00:00:00 2001 From: Mehdy Bohlool Date: Thu, 12 Nov 2026 20:35:17 +0000 Subject: [PATCH] [SDY][MaxDiffusion][cleanup] Remove obsolete jax_use_shardy_partitioner=False updates in attention_flax.py `jax_use_shardy_partitioner` is now a deprecated no-op in JAX because Shardy is always enabled. --- src/maxdiffusion/models/attention_flax.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/src/maxdiffusion/models/attention_flax.py b/src/maxdiffusion/models/attention_flax.py index d1d9937e3..a7e0fca7d 100644 --- a/src/maxdiffusion/models/attention_flax.py +++ b/src/maxdiffusion/models/attention_flax.py @@ -3324,8 +3324,6 @@ def __init__( if attention_kernel == "cudnn_flash_te": from transformer_engine.jax.flax.transformer import DotProductAttention # pytype: disable=import-error - jax.config.update("jax_use_shardy_partitioner", False) - dpa_layer = DotProductAttention( head_dim=dim_head, num_attention_heads=heads, @@ -3437,8 +3435,6 @@ def setup(self): if self.attention_kernel == "cudnn_flash_te": from transformer_engine.jax.flax.transformer import DotProductAttention # pytype: disable=import-error - jax.config.update("jax_use_shardy_partitioner", False) - dpa_layer = DotProductAttention( head_dim=self.dim_head, num_attention_heads=self.heads,