From fa8b9eb302e87ec4bd9a3a0794bf7a8898eb1f07 Mon Sep 17 00:00:00 2001
From: dfridovi <dfk@utexas.edu>
Date: Tue, 26 Nov 2024 08:21:39 -0600
Subject: [PATCH] reverse mode works except for handling ZeroTangents

---
 src/AutoDiff.jl  | 22 ++++++++++------------
 src/game.jl      |  2 +-
 src/solver.jl    |  4 ++--
 test/runtests.jl | 15 +++++++++------
 4 files changed, 22 insertions(+), 21 deletions(-)

diff --git a/src/AutoDiff.jl b/src/AutoDiff.jl
index 3e71e14..0372c71 100644
--- a/src/AutoDiff.jl
+++ b/src/AutoDiff.jl
@@ -25,19 +25,16 @@ function _solve_jacobian_θ(mcp::MCPSolver.PrimalDualMCP, solution, θ)
     )
 
     (; x, y, s, ϵ) = solution
-    ∂z∂θ = -mcp.∇F_z(x, y, s; θ, ϵ) \ mcp.∇F_θ(x, y, s; θ, ϵ)
+    ∂z∂θ =
+        LinearAlgebra.qr(-collect(mcp.∇F_z(x, y, s; θ, ϵ)), LinearAlgebra.ColumnNorm()) \
+        collect(mcp.∇F_θ(x, y, s; θ, ϵ))
 
-    SparseArrays.sparse(∂z∂θ)
+    ∂z∂θ
 end
 
-function ChainRulesCore.rrule(
-    ::typeof(MCPSolver.solve),
-    solver_type::MCPSolver.SolverType,
-    mcp::MCPSolver.PrimalDualMCP;
-    θ,
-    kwargs...,
-)
-    solution = MCPSolver.solve(solver_type, mcp; θ, kwargs...)
+function ChainRulesCore.rrule(::typeof(MCPSolver.solve), solver_type, mcp, θ; kwargs...)
+    println("yoyoyyo")
+    solution = MCPSolver.solve(solver_type, mcp, θ; kwargs...)
     project_to_θ = ChainRulesCore.ProjectTo(θ)
 
     function solve_pullback(∂solution)
@@ -52,6 +49,7 @@ function ChainRulesCore.rrule(
             ∂l∂x = ∂solution.x
             ∂l∂y = ∂solution.y
             ∂l∂s = ∂solution.s
+#            @infiltrate
             project_to_θ(∂z∂θ' * [∂l∂x; ∂l∂y; ∂l∂s])
         end
 
@@ -63,8 +61,8 @@ end
 
 function MCPSolver.solve(
     solver_type::MCPSolver.SolverType,
-    mcp::MCPSolver.PrimalDualMCP;
-    θ::AbstractVector{<:ForwardDiff.Dual{T}},
+    mcp::MCPSolver.PrimalDualMCP,
+    θ::AbstractVector{<:ForwardDiff.Dual{T}};
     kwargs...,
 ) where {T}
     # strip off the duals
diff --git a/src/game.jl b/src/game.jl
index 69534ef..86a203d 100644
--- a/src/game.jl
+++ b/src/game.jl
@@ -156,7 +156,7 @@ function solve(
     y₀ = ones(sum(game.dims.μ) + game.dims.μ̃),
     tol = 1e-4,
 )
-    (; x, y, s, kkt_error, status) = solve(solver_type, game.mcp; θ, x₀, y₀, tol)
+    (; x, y, s, kkt_error, status) = solve(solver_type, game.mcp, θ; x₀, y₀, tol)
 
     # Unpack primals per-player for ease of access later.
     end_dims = cumsum(game.dims.x)
diff --git a/src/solver.jl b/src/solver.jl
index 069d2e9..1ebb2ee 100644
--- a/src/solver.jl
+++ b/src/solver.jl
@@ -17,8 +17,8 @@ when the previous subproblem is solved in fewer iterations.
 """
 function solve(
     ::InteriorPoint,
-    mcp::PrimalDualMCP;
-    θ,
+    mcp::PrimalDualMCP,
+    θ;
     x₀ = zeros(mcp.unconstrained_dimension),
     y₀ = ones(mcp.constrained_dimension),
     tol = 1e-4,
diff --git a/test/runtests.jl b/test/runtests.jl
index cb131c3..136f673 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -45,7 +45,7 @@ using FiniteDiff: FiniteDiff
             constrained_dimension = length(b),
             parameter_dimension = size(M, 1),
         )
-        sol = MCPSolver.solve(MCPSolver.InteriorPoint(), mcp; θ)
+        sol = MCPSolver.solve(MCPSolver.InteriorPoint(), mcp, θ)
 
         check_solution(sol)
     end
@@ -57,7 +57,7 @@ using FiniteDiff: FiniteDiff
             fill(Inf, size(M, 1) + length(b));
             parameter_dimension = size(M, 1),
         )
-        sol = MCPSolver.solve(MCPSolver.InteriorPoint(), mcp; θ)
+        sol = MCPSolver.solve(MCPSolver.InteriorPoint(), mcp, θ)
 
         check_solution(sol)
     end
@@ -73,15 +73,18 @@ using FiniteDiff: FiniteDiff
         )
 
         function f(θ)
-            sol = MCPSolver.solve(MCPSolver.InteriorPoint(), mcp; θ)
-            sum(sol.x .^ 2) + sum(sol.y .^ 2)
+            sol = MCPSolver.solve(MCPSolver.InteriorPoint(), mcp, θ)
+            #@infiltrate
+            sum(sol.x .^ 2) + sum(sol.y .^ 2) + sum(sol.s .^ 2)
         end
 
+        #@infiltrate
+
         ∇_autodiff_reverse = only(Zygote.gradient(f, θ))
-        ∇_autodiff_forward = only(Zygote.gradient(θ -> Zygote.forwarddiff(f, θ), θ))
+        #∇_autodiff_forward = only(Zygote.gradient(θ -> Zygote.forwarddiff(f, θ), θ))
         ∇_finitediff = FiniteDiff.finite_difference_gradient(f, θ)
         @test isapprox(∇_autodiff_reverse, ∇_finitediff; atol = 1e-3)
-        @test isapprox(∇_autodiff_reverse, ∇_autodiff_forward; atol = 1e-3)
+        #@test isapprox(∇_autodiff_reverse, ∇_autodiff_forward; atol = 1e-3)
     end
 end