Fix .inst -> .stmt field renaming. Uncomment tests that have been bro…

…ken previously
dfdx · Mar 20, 2024 · b29bea6 · b29bea6
1 parent a2e2773
commit b29bea6
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 11 deletions.
diff --git a/src/trace.jl b/src/trace.jl
@@ -17,6 +17,17 @@ const VecOrTuple = Union{Tuple, Vector}
 #                                  Frame                                      #
 ###############################################################################
 
+# see https://github.com/dfdx/Umlaut.jl/issues/58
+function new_exprs(ir::IRCode)
+    stmts = ir.new_nodes.stmts
+    if hasfield(typeof(stmts), :inst)
+        return stmts.inst
+    else
+        return stmts.stmt
+    end
+end
+
+
 """
     block_expressions(ir::IRCode)
 
@@ -25,7 +36,7 @@ Returns Vector{block_info}, where block_info is Vector{ssa_id => expr}
 """
 function block_expressions(ir::IRCode)
     # new statements
-    new_exs = ir.new_nodes.stmts.inst
+    new_exs = new_exprs(ir)
     # where to insert them
     new_positions = [(info.attach_after ? info.pos + 1 : info.pos)
                     for info in ir.new_nodes.info]

diff --git a/test/test_trace.jl b/test/test_trace.jl
@@ -732,11 +732,11 @@ function conditionally_defined_tester(x)
 end
 
 ## broken in v1.10
-# @testset "undef in PhiNode" begin
-#     res, tape = trace(conditionally_defined_tester, 5.0)
-#     @test res == conditionally_defined_tester(5.0)
-#     @test play!(tape, conditionally_defined_tester, 5.0) == res
-# end
+@testset "undef in PhiNode" begin
+    res, tape = trace(conditionally_defined_tester, 5.0)
+    @test res == conditionally_defined_tester(5.0)
+    @test play!(tape, conditionally_defined_tester, 5.0) == res
+end
 
 ###############################################################################
 
@@ -748,8 +748,8 @@ function undefcheck_tester(x)
 end
 
 # broken in v1.10
-# @testset "undefcheck" begin
-#     res, tape = trace(undefcheck_tester, 5.0)
-#     @test res == undefcheck_tester(5.0)
-#     @test play!(tape, undefcheck_tester, 5.0) == res
-# end
+@testset "undefcheck" begin
+    res, tape = trace(undefcheck_tester, 5.0)
+    @test res == undefcheck_tester(5.0)
+    @test play!(tape, undefcheck_tester, 5.0) == res
+end