Merge pull request #112 from ACEsuit/co/bugs

Bugfixes

Author: cortner

src/interface.jl

@@ -80,10 +80,11 @@ _evaluate!(P, dP, basis::AbstractP4MLBasis, X) =
 # But ka_evaluate! can also be called with CPU arrays to enable testing 
 # KA kernels also on the CPU. 
 
-evaluate!(P::AbstractGPUArray, basis::AbstractP4MLBasis, x::BATCH) = 
+evaluate!(P::AbstractGPUArray, basis::AbstractP4MLBasis, x::BATCH, args...) = 
 		ka_evaluate!(P, basis, x)
 
-evaluate_ed!(P::AbstractGPUArray, dP::AbstractGPUArray, basis::AbstractP4MLBasis, x::BATCH) = 
+evaluate_ed!(P::AbstractGPUArray, dP::AbstractGPUArray, 
+             basis::AbstractP4MLBasis, x::BATCH, args...) = 
 		ka_evaluate_ed!(P, dP, basis, x)      
 
 function ka_evaluate!(P, basis::AbstractP4MLBasis, x::BATCH) 
@@ -194,12 +195,12 @@ function _with_safe_alloc(fcall, args...)
    return fcall(outputs..., args...)
 end
 
-function _with_safe_alloc(fcall, basis::AbstractP4MLBasis, X::BATCH) 
+function _with_safe_alloc(fcall, basis::AbstractP4MLBasis, X::BATCH, args...) 
    _alczero(T, args...) = fill!( similar(X, T, args...), zero(T) )
 
-   allocinfo = _tup_whatalloc(fcall, basis, X)
+   allocinfo = _tup_whatalloc(fcall, basis, X, args...)
    outputs = ntuple(i -> _alczero(allocinfo[i]...), length(allocinfo))
-   return fcall(outputs..., basis, X)
+   return fcall(outputs..., basis, X, args...)
 end
 
 # --------------------------------------- 

src/transformed.jl

@@ -87,12 +87,19 @@ _valtype(::Nothing, T::Type) = T
 # CPU SIMD kernel 
 # 
 
+_evaluate!(P, dP, tbasis::TransformedBasis, x::AbstractVector, 
+           ps::Nothing, st::Nothing) =
+   _evaluate!(P, dP, tbasis, x, 
+             (trans = NamedTuple(), basis = NamedTuple()), 
+             (trans = NamedTuple(), basis = NamedTuple()), )
 
-function _evaluate!(P, dP::Nothing, 
+
+function _evaluate!(P, dP, 
                     tbasis::TransformedBasis,
                     x::AbstractVector, 
                     ps, st)
    nX = length(x)
+   @assert isnothing(dP) "_evaluate! not implemented for dP != nothing"
 
    @no_escape begin
       # [1] Stage 1 - transform the inputs 

test/test_lux.jl

@@ -68,8 +68,7 @@ test_bases = [ chebyshev_basis(10),
 for basis in test_bases
 # basis = test_bases[12]
    @info("Lux layer test for $(typeof(basis).name.name)")
-   local B1, B2, x
-   local ps, st
+   local B1, B2, x, X, ps, st
    nX = rand(8:16)
    X = [ _generate_input(basis) for _ = 1:nX ]
    B1 = basis(X)