From 1daeb5956d026dd4b1a5790ff35385d39a6a5860 Mon Sep 17 00:00:00 2001
From: Avik Pal <avikpal@mit.edu>
Date: Tue, 9 Apr 2024 12:28:49 -0400
Subject: [PATCH] Test if ChainRules problem is resolved

---
 Project.toml            |  6 +-----
 ext/LuxChainRulesExt.jl | 18 ------------------
 2 files changed, 1 insertion(+), 23 deletions(-)
 delete mode 100644 ext/LuxChainRulesExt.jl

diff --git a/Project.toml b/Project.toml
index b6da5a9a6..4e6be37a8 100644
--- a/Project.toml
+++ b/Project.toml
@@ -1,7 +1,7 @@
 name = "Lux"
 uuid = "b2108857-7c20-44ae-9111-449ecde12c47"
 authors = ["Avik Pal <avikpal@mit.edu> and contributors"]
-version = "0.5.33"
+version = "0.5.34"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
@@ -29,7 +29,6 @@ WeightInitializers = "d49dbf32-c5c2-4618-8acc-27bb2598ef2d"
 
 [weakdeps]
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
-ChainRules = "082447d4-558c-5d27-93f4-14fc19e9eca2"
 ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"
 Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
 LuxAMDGPU = "83120cb1-ca15-4f04-bf3b-6967d2e6b60b"
@@ -43,7 +42,6 @@ Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [extensions]
-LuxChainRulesExt = "ChainRules"
 LuxComponentArraysExt = "ComponentArrays"
 LuxComponentArraysReverseDiffExt = ["ComponentArrays", "ReverseDiff"]
 LuxFluxExt = "Flux"
@@ -63,7 +61,6 @@ Adapt = "4"
 Aqua = "0.8.4"
 ArrayInterface = "7.8"
 CUDA = "5.2"
-ChainRules = "1.62"
 ChainRulesCore = "1.21"
 ComponentArrays = "0.15.11"
 ConcreteStructs = "0.2.3"
@@ -108,7 +105,6 @@ julia = "1.10"
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
-ChainRules = "082447d4-558c-5d27-93f4-14fc19e9eca2"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"
 ExplicitImports = "7d51a73a-1435-4ff3-83d9-f097790105c7"
diff --git a/ext/LuxChainRulesExt.jl b/ext/LuxChainRulesExt.jl
deleted file mode 100644
index 9ecf874d0..000000000
--- a/ext/LuxChainRulesExt.jl
+++ /dev/null
@@ -1,18 +0,0 @@
-module LuxChainRulesExt
-
-using ChainRules: ChainRules
-
-# https://github.com/FluxML/Zygote.jl/pull/1328 broke the RNNs completely. Putting an
-# emergency patch here
-function ChainRules._setindex_zero(
-        x::Vector{<:AbstractArray{T}}, dy, inds::Integer...) where {T <: Number}
-    return [fill!(similar(xᵢ), 0) for xᵢ in x]
-end
-
-function ChainRules.∇getindex!(
-        dx::Vector{<:AbstractArray{T}}, dy, inds::Integer...) where {T <: Number}
-    dx[inds...] .+= dy
-    return dx
-end
-
-end