From 50e8d9c2b1b682a800bff2908db6a39629e5b9cc Mon Sep 17 00:00:00 2001 From: Will Tebbutt Date: Sun, 20 Oct 2024 13:22:55 +0100 Subject: [PATCH] Make use of new MistyClosure functionality (#296) --- Project.toml | 4 ++-- src/interpreter/s2s_reverse_mode_ad.jl | 10 ++-------- 2 files changed, 4 insertions(+), 10 deletions(-) diff --git a/Project.toml b/Project.toml index ae8b6c941..557f7bea8 100644 --- a/Project.toml +++ b/Project.toml @@ -1,7 +1,7 @@ name = "Mooncake" uuid = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6" authors = ["Will Tebbutt, Hong Ge, and contributors"] -version = "0.4.12" +version = "0.4.13" [deps] ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" @@ -55,7 +55,7 @@ Graphs = "1" JET = "0.9" LogDensityProblemsAD = "1" LuxLib = "1.2 - 1.3.3" -MistyClosures = "1" +MistyClosures = "1.0.3" NNlib = "0.9" PDMats = "0.11" Setfield = "1" diff --git a/src/interpreter/s2s_reverse_mode_ad.jl b/src/interpreter/s2s_reverse_mode_ad.jl index aeb5498e7..8c5c30151 100644 --- a/src/interpreter/s2s_reverse_mode_ad.jl +++ b/src/interpreter/s2s_reverse_mode_ad.jl @@ -930,14 +930,8 @@ function build_rrule( # @show length(stmt(ir.stmts)) # @show length(stmt(optimised_fwds_ir.stmts)) # @show length(stmt(optimised_pb_ir.stmts)) - fwds_oc = MistyClosure( - OpaqueClosure(optimised_fwds_ir, shared_data...; do_compile=true), - optimised_fwds_ir, - ) - pb_oc = MistyClosure( - OpaqueClosure(optimised_pb_ir, shared_data...; do_compile=true), - optimised_pb_ir, - ) + fwds_oc = MistyClosure(optimised_fwds_ir, shared_data...; do_compile=true) + pb_oc = MistyClosure(optimised_pb_ir, shared_data...; do_compile=true) # Compute the signature. Needs careful handling with varargs. sig = sig_or_mi isa Core.MethodInstance ? sig_or_mi.specTypes : sig_or_mi