format

Author: AstitvaAggarwal

ext/MooncakeDiffEqBaseExt.jl

@@ -3,7 +3,6 @@ module MooncakeDiffEqBaseExt
 using DiffEqBase, Mooncake
 using DiffEqBase: SciMLBase
 using SciMLBase: ADOriginator, MooncakeOriginator
-println("mooncakke diff eqbase ext loaded !")
 import Mooncake:
     rrule!!,
     CoDual,
@@ -17,36 +16,6 @@ import Mooncake:
     MinimalCtx,
     NoPullback
 
-# @from_rrule(
-#     MinimalCtx,
-#     Tuple{
-#         typeof(DiffEqBase.solve_up),
-#         DiffEqBase.AbstractDEProblem,
-#         Union{Nothing,DiffEqBase.AbstractSensitivityAlgorithm},
-#         Any,
-#         Any,
-#         Any,
-#     },
-#     true,
-# )
-
-# # Dispatch for auto-alg
-# @from_rrule(
-#     MinimalCtx,
-#     Tuple{
-#         typeof(DiffEqBase.solve_up),
-#         DiffEqBase.AbstractDEProblem,
-#         Union{Nothing,DiffEqBase.AbstractSensitivityAlgorithm},
-#         Any,
-#         Any,
-#     },
-#     true,
-# )
-
-# @zero_adjoint MinimalCtx Tuple{typeof(DiffEqBase.numargs),Any}
-# @mooncake_overlay DiffEqBase.set_mooncakeoriginator_if_mooncake(x::ADOriginator) =
-#     MooncakeOriginator
-
 @is_primitive MinimalCtx Tuple{
     typeof(DiffEqBase.set_mooncakeoriginator_if_mooncake),SciMLBase.ChainRulesOriginator
 }
@@ -56,21 +25,4 @@ function rrule!!(
 )
     return zero_fcodual(SciMLBase.MooncakeOriginator()), NoPullback(f, X)
 end
-# o method matching rrule!!(::CoDual{typeof(DiffEqBase.set_mooncakeoriginator_if_mooncake), Mooncake.NoFData}, ::CoDual{SciMLBase.ChainRulesOriginator, Mooncake.NoFData})
-# Mooncake.DerivedRule{
-#     Tuple{
-#         typeof(DiffEqBase.set_mooncakeoriginator_if_mooncake),SciMLBase.ChainRulesOriginator
-#     },
-#     Tuple{
-#         Mooncake.CoDual{
-#             typeof(DiffEqBase.set_mooncakeoriginator_if_mooncake),Mooncake.NoFData
-#         },
-#         Mooncake.CoDual{SciMLBase.ChainRulesOriginator,Mooncake.NoFData},
-#     },
-#     Mooncake.CoDual{Type{SciMLBase.MooncakeOriginator},Mooncake.NoFData},
-#     Tuple{Mooncake.NoRData},
-#     Tuple{Mooncake.NoRData,Mooncake.NoRData},
-#     false,
-#     Val{2},
-# }
-end
+end

test/ext/diffeqbase/Project.toml

@@ -1,7 +1,6 @@
 [deps]
 AllocCheck = "9b6a8646-10ed-4001-bbdc-1d2f46dfbb1a"
 DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
-JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
 Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
 StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"

test/ext/diffeqbase/diffeqbase.jl

@@ -2,7 +2,7 @@ using Pkg
 Pkg.activate(@__DIR__)
 Pkg.develop(; path=joinpath(@__DIR__, "..", "..", ".."))
 
-using AllocCheck, JET
+using AllocCheck
 using Mooncake, StableRNGs, Test, DiffEqBase
 using DiffEqBase: SciMLBase
 using Mooncake.TestUtils: test_rule
@@ -12,11 +12,11 @@ using Mooncake.TestUtils: test_rule
         interface_only, perf_flag, is_primitive, f, fargs
     ) in [(
         false,
-        :stability_and_allocs,
+        :allocs,
         true,
         DiffEqBase.set_mooncakeoriginator_if_mooncake,
         SciMLBase.ChainRulesOriginator(),
     )]
         test_rule(StableRNG(123), f, fargs; interface_only, perf_flag, is_primitive)
     end
-end
+end