Merge pull request #79 from ACEsuit/chho/fix_releaseinput

Fix release input in `PooledSparseProduct` and tuple inputs

src/ace/sparseprodpool.jl

@@ -574,15 +574,28 @@ import LuxCore: AbstractExplicitLayer, initialparameters, initialstates
 
 struct PooledSparseProductLayer{NB} <: AbstractExplicitLayer 
    basis::PooledSparseProduct{NB}
+   meta::Dict{String, Any}
+   release_input::Bool
 end
 
-lux(basis::PooledSparseProduct) = PooledSparseProductLayer(basis)
+function lux(basis::PooledSparseProduct; 
+               name = String(nameof(typeof(basis))), 
+               meta = Dict{String, Any}("name" => name),
+               release_input = true)
+   @assert haskey(meta, "name")
+   return PooledSparseProductLayer(basis, meta, release_input)
+end
 
 initialparameters(rng::AbstractRNG, layer::PooledSparseProductLayer) = 
       NamedTuple() 
 
 initialstates(rng::AbstractRNG, layer::PooledSparseProductLayer) = 
       NamedTuple()
 
-(l::PooledSparseProductLayer)(BB, ps, st) = 
-      evaluate(l.basis, BB), st 
+(l::PooledSparseProductLayer)(BB, ps, st) = begin
+   out = evaluate(l.basis, BB)
+   if l.release_input
+      release!.(BB)
+   end
+   return out, st
+end

src/lux.jl

@@ -66,7 +66,7 @@ function evaluate(l::PolyLuxLayer, X, ps, st)
    out = acquire!(st.pool, _outsym(X), _out_size(l.basis, X), _valtype(l.basis, X))
    evaluate!(out, l.basis, X, ps)
    if l.release_input
-      release!(X)
+      release!.(X)
    end
    return out, st
 end

test/test_linear.jl

@@ -68,7 +68,8 @@ for (feat, in_size, out_fun) in zip(feature_arr, in_size_arr, out_fun_arr)
          print_tf(@test fdtest(F, dF, 0.0; verbose=false))
       end
    end
-
+   println()
+
    @info("Testing evaluate")
    for ntest = 1:30
       x = randn(in_size)