Fixed Zygote tests

Author: divital-coder

src/basic.jl

@@ -170,22 +170,19 @@ end
 # Out-of-place: v is action vector, u is update vector
 function (nn::NullOperator)(v::AbstractVecOrMat, u, p, t; kwargs...)
     @assert size(v, 1) == nn.len
-    update_coefficients(nn, u, p, t; kwargs...)
     zero(v)
 end
 
 # In-place: w is destination, v is action vector, u is update vector
 function (nn::NullOperator)(w::AbstractVecOrMat, v::AbstractVecOrMat, u, p, t; kwargs...)
     @assert size(v, 1) == nn.len
-    update_coefficients!(nn, u, p, t; kwargs...)
     lmul!(false, w)
     w
 end
 
 # In-place with scaling: w = α*(nn*v) + β*w
 function (nn::NullOperator)(w::AbstractVecOrMat, v::AbstractVecOrMat, u, p, t, α, β; kwargs...)
     @assert size(v, 1) == nn.len
-    update_coefficients!(nn, u, p, t; kwargs...)
     lmul!(β, w)
     w
 end

test/runtests.jl

@@ -17,6 +17,6 @@ using SafeTestsets
         include("total.jl")
     end
     @time @safetestset "Zygote.jl" begin
-        # include("zygote.jl")
+        include("zygote.jl")
     end
 end

test/zygote.jl

@@ -66,19 +66,26 @@ for (LType, L) in ((IdentityOperator, IdentityOperator(N)),
     (AddedScalarOperator, α + α),
     (ComposedScalarOperator, α * α))
     @assert L isa LType
+    
+    # Cache the operator for efficient application
+    L_cached = cache_operator(L, u0)
 
+    # Updated loss function using the new interface:
+    # v is the action vector, u0 is the update vector
     loss_mul = function (p)
         v = Diagonal(p) * u0
-        w = L(v, p, t)
+        # Use new interface: L(v, u, p, t)
+        w = L_cached(v, u0, p, t)
         l = sum(w)
     end
 
     loss_div = function (p)
         v = Diagonal(p) * u0
-
-        L = update_coefficients(L, v, p, t)
-        w = L \ v
-
+        
+        # Update coefficients first, then apply inverse
+        L_updated = update_coefficients(L_cached, u0, p, t)
+        w = L_updated \ v
+        
         l = sum(w)
     end
 
@@ -99,4 +106,4 @@ for (LType, L) in ((IdentityOperator, IdentityOperator(N)),
             @test !isa(g_div, Nothing)
         end
     end
-end
+end