add wan vae attention test

Author: jfacevedo-google

src/maxdiffusion/models/wan/autoencoder_kl_wan.py

@@ -364,9 +364,45 @@ def __init__(
       rngs: nnx.Rngs
   ):
     self.dim = dim
+    self.norm = WanRMS_norm(rngs=rngs, dim=dim, channel_first=False)
+    self.to_qkv = nnx.Conv(
+      in_features=dim,
+      out_features=dim * 3,
+      kernel_size=1,
+      rngs=rngs
+    )
+    self.proj = nnx.Conv(
+      in_features=dim,
+      out_features=dim,
+      kernel_size=1,
+      rngs=rngs
+    )
 
   def __call__(self, x: jax.Array):
-    return x
+    batch_size, time, height, width, channels = x.shape
+    identity = x
+    
+    x = x.reshape(batch_size * time, height, width, channels)
+    x = self.norm(x)
+
+    qkv = self.to_qkv(x) # Output: (N*D, H, W, C * 3)
+
+    qkv = qkv.reshape(batch_size*time, 1, channels * 3, -1)
+    qkv = jnp.transpose(qkv, (0, 1, 3, 2))
+    q, k, v = jnp.split(qkv, 3, axis=-1)
+
+    x = jax.nn.dot_product_attention(q, k, v)
+    x = jnp.squeeze(x, 1).reshape(batch_size * time, height, width, channels)
+
+    #output projection
+    x = self.proj(x)
+
+    # Reshape back
+    x = x.reshape(batch_size, time, height, width, channels)
+
+    return x + identity
+
+
 
 class WanMidBlock(nnx.Module):
   def __init__(

src/maxdiffusion/tests/wan_vae_test.py

@@ -33,7 +33,8 @@
   WanResidualBlock,
   WanRMS_norm,
   WanResample,
-  ZeroPaddedConv2D
+  ZeroPaddedConv2D,
+  WanAttentionBlock
 )
 
 CACHE_T = 2
@@ -322,7 +323,7 @@ def test_3d_conv(self):
   def test_wan_residual(self):
     key = jax.random.key(0)
     rngs = nnx.Rngs(key)
-    # one test
+    # --- Test Case 1: same in/out dim ---
     in_dim = out_dim = 96
     batch = 1
     t = 1
@@ -341,7 +342,7 @@ def test_wan_residual(self):
     dummy_output = wan_residual_block(dummy_input)
     assert dummy_output.shape == expected_output_shape
 
-    # another test
+    # --- Test Case 1: different in/out dim ---
     in_dim = 96
     out_dim = 196
     expected_output_shape = (batch, t, height, width, out_dim)
@@ -355,8 +356,22 @@ def test_wan_residual(self):
     dummy_output = wan_residual_block(dummy_input)
     assert dummy_output.shape == expected_output_shape
 
-
-
+  def test_wan_attention(self):
+    key = jax.random.key(0)
+    rngs = nnx.Rngs(key)
+    dim = 384
+    batch = 1
+    t = 1
+    height = 60
+    width = 90
+    input_shape=(batch, t, height, width, dim)
+    wan_attention = WanAttentionBlock(
+      dim=dim,
+      rngs=rngs
+    )
+    dummy_input = jnp.ones(input_shape)
+    output = wan_attention(dummy_input)
+    assert output.shape == input_shape
 
   def test_wan_encode(self):
     key = jax.random.key(0)