aoti diag: functional none-gate

app.py

@@ -483,8 +483,9 @@ def aoti_diag() -> str:
         return vsa_h3.sparse_attention(q, k, v, None if gate_h is None else gate_h.transpose(1, 2), geometry)
 
     def functional(gate_h):
+        gate_t = None if gate_h is None else gate_h.transpose(1, 2)
         return aoti_attention.sparse_attention_functional(
-            q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), gate_h.transpose(1, 2),
+            q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), gate_t,
             geometry.untile_index, geometry.variable_block_sizes, geometry.tile_divisor,
             geometry.topk, geometry.num_prefix_tiles,
         )