From a276da53128d1b02caa681adc8d8edd5f0d7567d Mon Sep 17 00:00:00 2001
From: cheukhinhojerry <cheukhinhojerry@gmail.com>
Date: Fri, 14 Jun 2024 21:50:01 -0700
Subject: [PATCH] fix tests

---
 src/Polynomials4ML.jl |  2 +-
 src/linear.jl         |  2 +-
 test/test_linear.jl   | 17 ++++++-----------
 3 files changed, 8 insertions(+), 13 deletions(-)

diff --git a/src/Polynomials4ML.jl b/src/Polynomials4ML.jl
index d420840..d2ef943 100644
--- a/src/Polynomials4ML.jl
+++ b/src/Polynomials4ML.jl
@@ -72,7 +72,7 @@ include("sparseproduct.jl")
 # LinearLayer implementation
 # this is needed to better play with cached arrays + to give the correct 
 # behaviour when the feature dimension is different from expected. 
-# include("linear.jl")
+include("linear.jl")
 
 # generic machinery for wrapping poly4ml bases into lux layers 
 include("lux.jl")
diff --git a/src/linear.jl b/src/linear.jl
index 071dc45..77580e8 100644
--- a/src/linear.jl
+++ b/src/linear.jl
@@ -33,7 +33,7 @@ out, st = l(x, ps, st)
 println(out == x * transpose(W))) # true
 ```
 """
-struct LinearLayer{FEATFIRST} <: AbstractP4MLLayer
+struct LinearLayer{FEATFIRST} <: AbstractP4MLTensor
    in_dim::Integer
    out_dim::Integer
    @reqfields()
diff --git a/test/test_linear.jl b/test/test_linear.jl
index 56d2871..61a43be 100644
--- a/test/test_linear.jl
+++ b/test/test_linear.jl
@@ -34,7 +34,7 @@ for (feat, in_size, out_fun) in zip(feature_arr, in_size_arr, out_fun_arr)
    if !feat 
       @info("Testing evaluate on vector input vs batch input")
       for ntest = 1:30 
-         X = randn(N, in_d) 
+         X = randn(N, in_d)
          Y1, _ = l(X, ps, st)
          Y2 = hcat([l(X[i,:], ps, st)[1] for i = 1:N]...)'
          Y3 = hcat([ps.W * X[i,:] for i = 1:N]...)'
@@ -52,7 +52,7 @@ for (feat, in_size, out_fun) in zip(feature_arr, in_size_arr, out_fun_arr)
          u = randn(size(val))
          F(t) = dot(u, l(_BB(t), ps, st)[1])
          dF(t) = begin
-            val, pb = Zygote.pullback(LuxCore.apply, l, _BB(t), ps, st)
+            val, pb = Zygote.pullback(P4ML.evaluate, l, _BB(t), ps, st)
             ∂BB = pb((u, st))[2]
             return dot(∂BB, bu)
          end
@@ -79,7 +79,7 @@ for (feat, in_size, out_fun) in zip(feature_arr, in_size_arr, out_fun_arr)
       u = randn(size(val))
       F(t) = dot(u, l(_BB(t), ps, st)[1])
       dF(t) = begin
-         val, pb = Zygote.pullback(LuxCore.apply, l, _BB(t), ps, st)
+         val, pb = Zygote.pullback(P4ML.evaluate, l, _BB(t), ps, st)
          ∂BB = pb((u, st))[2]
          return dot(∂BB, bu)
       end
@@ -99,7 +99,7 @@ for (feat, in_size, out_fun) in zip(feature_arr, in_size_arr, out_fun_arr)
       u = randn(size(val))
       F(t) = dot(u, l(x, re([_BB(t)...]), st)[1])
       dF(t) = begin
-         val, pb = Zygote.pullback(LuxCore.apply, l, x, re([_BB(t)...]), st)
+         val, pb = Zygote.pullback(P4ML.evaluate, l, x, re([_BB(t)...]), st)
          ∂BB = pb((u, st))[3]
          return dot(∂BB[1], bu)
       end
@@ -113,21 +113,16 @@ for (feat, in_size, out_fun) in zip(feature_arr, in_size_arr, out_fun_arr)
    # test_rrule(LuxCore.apply, l, x, ps, st)
 end
 
-##
-
+#
 
 # check which matmul it is calling
 # l = P4ML.LinearLayer(in_d, out_d; feature_first = false)
 # ps, st = LuxCore.setup(MersenneTwister(1234), l)
 # X = rand(N, in_d)
-# using ObjectPools
-# release!(X)
-# X = rand(N,in_d)
 
 # @profview let l = l, ps = ps, st = st, X = X
 #    for _ = 1:100_000
-#       out = l(X, ps, st)[1]
-#       release!(out)
+#       l(X, ps, st)
 #    end
 # end