Fix reciprocal op in TVM relay PyTorch frontend

tenstorrent · Sep 10, 2024 · 999580e · 999580e
1 parent 22a23e8
commit 999580e
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 5 deletions.
diff --git a/forge/test/mlir/test_ops.py b/forge/test/mlir/test_ops.py
@@ -354,7 +354,7 @@ def forward(self, x):
 @pytest.mark.parametrize("token_num", [12])
 @pytest.mark.parametrize("embedding_dim", [3200])
 def test_embedding(vocab_size, token_num, embedding_dim):
-    compiler_cfg = pyforge.config._get_global_compiler_config()
+    compiler_cfg = forge.config._get_global_compiler_config()
     compiler_cfg.enable_tvm_cpu_fallback = False
 
     class Embedding(nn.Module):
@@ -372,14 +372,13 @@ def forward(self, x):
     framework_model = Embedding()
     fw_out = framework_model(*inputs)
 
-    compiled_model = pyforge.compile(framework_model, sample_inputs=inputs)
+    compiled_model = forge.compile(framework_model, sample_inputs=inputs)
     co_out = compiled_model(*inputs)
 
     co_out = [co.to("cpu") for co in co_out]
     assert compare_with_golden_pcc(golden=fw_out, calculated=co_out[0], pcc=0.99)
 
 
-@pytest.mark.xfail(reason="Program expects 2 inputs, found 1 in input tensors vector")
 @pytest.mark.parametrize("shape", [
     (7,),           # 1D tensor
     (32,),          # 1D tensor
@@ -405,7 +404,7 @@ def forward(self, x):
     framework_model = Reciprocal()
     fw_out = framework_model(*inputs)
 
-    compiled_model = pyforge.compile(framework_model, sample_inputs=inputs)
+    compiled_model = forge.compile(framework_model, sample_inputs=inputs)
     co_out = compiled_model(*inputs)
 
     co_out = [co.to("cpu") for co in co_out]

diff --git a/third_party/tvm b/third_party/tvm