From 4638239e1c925307b87754f1bd09dc0daf19141b Mon Sep 17 00:00:00 2001
From: Brian Chen <ToucheSir@users.noreply.github.com>
Date: Sun, 7 Jan 2024 21:55:01 -0800
Subject: [PATCH] Add rule for `push!(::IdSet)`

This is required on 1.11 because the underlying implementation has changed.
---
 src/lib/base.jl | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/src/lib/base.jl b/src/lib/base.jl
index c0efa83cf..3e3dc48ed 100644
--- a/src/lib/base.jl
+++ b/src/lib/base.jl
@@ -17,6 +17,16 @@ end
   end
 end
 
+# IdSet (needed for nested AD with implicit params)
+function _pullback(cx::AContext, ::typeof(push!), s::IdSet, @nospecialize(x))
+  res = push!(s, x)
+  function idset_push!_pullback(_)
+    Δ = pop!(grad_mut(cx, d), x, nothing)
+    (nothing, Δ, nothing)
+  end
+  return res, idset_push!_pullback
+end
+
 # Dictionaries
 
 grad_mut(d::AbstractDict) = Dict()