transpose the output

mshr-h · mshr-h · commit f6f6c1aeb9ce · 2024-09-17T23:49:27.000+09:00
diff --git a/python/tvm/relax/frontend/torch/fx_translator.py b/python/tvm/relax/frontend/torch/fx_translator.py
@@ -1015,7 +1015,9 @@ def _scaled_dot_product_attention(self, node: fx.Node) -> relax.Var:
             assert "float" in attn_mask.struct_info.dtype, msg
 
         return self.block_builder.emit(
-            relax.op.nn.attention(query, key, value, bias=attn_mask, causal_mask=causal_mask)
+            transpose_S_H(
+                relax.op.nn.attention(query, key, value, bias=attn_mask, causal_mask=causal_mask)
+            )
         )
 
     def _unbind(self, node: fx.Node) -> relax.Var:

Original file line number	Diff line number	Diff line change
`@@ -1015,7 +1015,9 @@ def _scaled_dot_product_attention(self, node: fx.Node) -> relax.Var:`
`1015`	`1015`	`assert "float" in attn_mask.struct_info.dtype, msg`
`1016`	`1016`
`1017`	`1017`	`return self.block_builder.emit(`
`1018`		`- relax.op.nn.attention(query, key, value, bias=attn_mask, causal_mask=causal_mask)`
	`1018`	`+ transpose_S_H(`
	`1019`	`+ relax.op.nn.attention(query, key, value, bias=attn_mask, causal_mask=causal_mask)`
	`1020`	`+ )`
`1019`	`1021`	`)`
`1020`	`1022`
`1021`	`1023`	`def _unbind(self, node: fx.Node) -> relax.Var:`