Merge pull request #110 from ACEsuit/lux

Make all bases Lux layers

Author: cortner

Project.toml

@@ -25,7 +25,7 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 WithAlloc = "fb1aa66a-603c-4c1d-9bc4-66947c7b08dd"
 
 [compat]
-ACEbase = "0.4.2"
+ACEbase = "0.4.5"
 BenchmarkTools = "1"
 Bumper = "0.7.0"
 ChainRulesCore = "1"

src/Polynomials4ML.jl

@@ -2,30 +2,35 @@ module Polynomials4ML
 
 # -------------- import ACEbase, Bumper, WithAlloc, Lux and related
 
-import ACEbase
+import ACEbase 
+
 import ACEbase: evaluate, evaluate_d, evaluate_ed, 
-                evaluate!, evaluate_d!, evaluate_ed!
+                evaluate!, evaluate_d!, evaluate_ed!, 
+                pullback, pullback!, pushforward, pushforward!
+
+import ChainRulesCore: rrule, frule, NoTangent, ZeroTangent
+import LuxCore: AbstractLuxLayer, initialparameters, initialstates                 
 
 using Bumper, WithAlloc
 import WithAlloc: whatalloc 
 
 using KernelAbstractions, GPUArraysCore
 
-using LuxCore, Random, StaticArrays
-import ChainRulesCore: rrule, frule, NoTangent, ZeroTangent
+using LuxCore, Random, StaticArrays, ChainRulesCore
 using ForwardDiff: Dual, extract_derivative
 using StaticArrays
-import LuxCore: AbstractLuxLayer, initialparameters, initialstates                 
+
 
 using Random: AbstractRNG   
 
 """
 `abstract type AbstractP4MLBasis end`
 
 Annotates types that map a low-dimensional input, scalar or `SVector`,
-to a vector of scalars (feature vector, embedding, basis...). 
+to a vector of scalars (feature vector, embedding, basis...). Can be used 
+as a `Lux` layer. 
 """
-abstract type AbstractP4MLBasis end
+abstract type AbstractP4MLBasis <: AbstractLuxLayer end
 
 
 """
@@ -64,16 +69,6 @@ Chebyshev polynomials, `index(basis, n)` returns `n+1`.
 function index end
 
 function orthpolybasis end
-function degree end 
-
-function pullback! end
-function pullback end
-function pushforward end
-function pushforward! end
-
-# some stuff to allow bases to overload some lux functionality ... 
-# how much of this should go into ACEbase? 
-function lux end 
 
 export orthpolybasis
 
@@ -115,9 +110,6 @@ include("atomicorbitals/atomicorbitals.jl")
 # RETIRE - to be discussed? 
 # include("linear.jl")
 
-# generic machinery for wrapping poly4ml bases into lux layers 
-include("lux.jl")
-
 # some nice utility functions to generate basis sets and other things  
 include("utils/utils.jl")
 

src/atomicorbitals/atomicorbitals.jl

@@ -139,8 +139,8 @@ function pullback_ps(∂Rnl, basis::AtomicOrbitals, X::AbstractVector{<: SVector
     map!(norm, R, X)
 
     # Rnl = output of evaluate(basis, X, ...)
-    Pn = evaluate(basis.Pn, X, ps.Pn, st.Pn)
-    Dn = evaluate(basis.Dn, X, ps.Dn, st.Dn)
+    Pn = evaluate(basis.Pn, R, ps.Pn, st.Pn)
+    Dn = evaluate(basis.Dn, R, ps.Dn, st.Dn)
     Ylm = evaluate(basis.Ylm, X, ps.Ylm, st.Ylm)
     ∂Pn = zeros(T, size(Pn))
     ∂Dn = zeros(T, size(Dn))
@@ -156,8 +156,8 @@ function pullback_ps(∂Rnl, basis::AtomicOrbitals, X::AbstractVector{<: SVector
         end 
     end 
 
-    ∂p_Pn = pullback_ps(∂Pn, basis.Pn, X, ps.Pn, st.Pn)
-    ∂p_Dn = pullback_ps(∂Dn, basis.Dn, X, ps.Dn, st.Dn)
+    ∂p_Pn = pullback_ps(∂Pn, basis.Pn, R, ps.Pn, st.Pn)
+    ∂p_Dn = pullback_ps(∂Dn, basis.Dn, R, ps.Dn, st.Dn)
     ∂p_Ylm = pullback_ps(∂Ylm, basis.Ylm, X, ps.Ylm, st.Ylm)
     return (Pn = ∂p_Pn, Dn = ∂p_Dn, Ylm = ∂p_Ylm)
 end

src/bernstein.jl

@@ -18,7 +18,8 @@ _generate_input(basis::BernsteinBasis) = rand()
 end
 
 
-function _evaluate!(P, dP, basis::BernsteinBasis{N}, x::AbstractVector{<:Real}) where {N}
+function _evaluate!(P, dP, basis::BernsteinBasis{N}, x::AbstractVector{<:Real}, 
+                     ps, st) where {N}
 
    n = N - 1                     
    WITHGRAD = !isnothing(dP)

src/interface.jl

@@ -217,3 +217,26 @@ evaluate_ed(l::AbstractP4MLBasis, args...) =
 evaluate_d(l::AbstractP4MLBasis, args...) = 
       evaluate_ed(l, args...)[2] 
 
+
+# ------------------------------------------------------------
+# Lux interface
+
+
+"""
+a fall-back method for `initalparameters` that all AbstractP4MLBasis
+should overload 
+"""
+_init_luxparams(rng::AbstractRNG, l::Any) = _init_luxparams(l)
+_init_luxparams(l) = NamedTuple() 
+
+_init_luxstate(rng::AbstractRNG, l::Any) = _init_luxstate(l)
+_init_luxstate(l) = NamedTuple() 
+
+initialparameters(rng::AbstractRNG, l::AbstractP4MLBasis) = 
+      _init_luxparams(rng, l)
+
+initialstates(rng::AbstractRNG, l::AbstractP4MLBasis) = 
+      _init_luxstate(rng, l)
+
+(l::AbstractP4MLBasis)(X, ps::NamedTuple, st::NamedTuple) = 
+      evaluate(l, X, ps, st), st 

src/lux.jl

@@ -48,11 +48,11 @@ _generate_input(basis::TransformedBasis{Nothing}) =
 # --------------------------------------------------------- 
 # Lux stuff
 
-_init_luxparams(rng::AbstractRNG, l::TransformedBasis) = 
+initialparameters(rng::AbstractRNG, l::TransformedBasis) = 
       ( trans = _init_luxparams(rng, l.trans),
         basis = _init_luxparams(rng, l.basis), )
 
-_init_luxstate(rng::AbstractRNG, l::TransformedBasis) =
+initialstates(rng::AbstractRNG, l::TransformedBasis) =
       ( trans = _init_luxparams(rng, l.trans),
         basis = _init_luxparams(rng, l.basis),  )
 
@@ -96,7 +96,7 @@ function _evaluate!(P, dP::Nothing,
 
    @no_escape begin
       # [1] Stage 1 - transform the inputs 
-      z1 = evaluate(tbasis.trans, x[1], ps.trans, st.trans) 
+      z1 = evaluate(tbasis.trans, x[1], ps.trans, st.trans)
       TZ = typeof(z1)
       Z = @alloc(TZ, nX) 
       @inbounds begin

src/transformed.jl

@@ -1,6 +1,7 @@
 using Polynomials4ML
 using Test
 
+##
 @testset "Polynomials4ML.jl" begin 
 
     # 1D Polynomials 
@@ -19,13 +20,14 @@ using Test
 
     # Misc
     @testset "Static Prod" begin include("test_staticprod.jl"); end
-    # @testset "Lux" begin include("test_lux.jl"); end 
 
     # Transformations 
     @testset "Transformed Basis" begin include("test_transformed.jl"); end
 
+    # Test lux interface 
+    @testset "Lux" begin include("test_lux.jl"); end 
+
     # TODO: restructure or move?? 
     # @testset "Sparse Product" begin include("test_sparseproduct.jl"); end 
     # @testset "Linear layer" begin include("test_linear.jl"); end 
-
 end

test/runtests.jl

@@ -14,7 +14,7 @@ basis = CTrigBasis(N)
 @info("      correctness")
 mm = natural_indices(basis)
 for ntest = 1:10
-   local x 
+   local x, P  
    x = 2*π * rand()
    P = basis(x)
    P_ref = [ exp(im * m.n * x) for m in mm ]

test/test_ctrig.jl

@@ -1,5 +1,5 @@
 using Polynomials4ML, Test, StaticArrays, LuxCore, Zygote, ForwardDiff
-using Polynomials4ML: lux, _generate_input
+using Polynomials4ML: _generate_input
 using Random: default_rng
 using ACEbase.Testing: println_slim, print_tf
 using LinearAlgebra: dot, I 
@@ -66,16 +66,15 @@ test_bases = [ chebyshev_basis(10),
 ##
 
 for basis in test_bases
-# basis = test_bases[13]
+# basis = test_bases[12]
    @info("Lux layer test for $(typeof(basis).name.name)")
    local B1, B2, x
    local ps, st
    nX = rand(8:16)
    X = [ _generate_input(basis) for _ = 1:nX ]
    B1 = basis(X)
-   l = lux(basis)
-   ps, st = LuxCore.setup(rng, l)
-   B2, _ = l(X, ps, st)
+   ps, st = LuxCore.setup(rng, basis)
+   B2, _ = basis(X, ps, st)
    println_slim(@test B1 == B2)
 
    if !( eltype(eltype(B1)) <: Real)
@@ -84,7 +83,7 @@ for basis in test_bases
    end
 
    # evaluate the basis and get the pullback operator  
-   val1, pb1 = Zygote.pullback(LuxCore.apply, l, X, ps, st)
+   val1, pb1 = Zygote.pullback(LuxCore.apply, basis, X, ps, st)
    val2, pb2 = Zygote.pullback(Polynomials4ML.evaluate, basis, X, ps, st)
    # evaluate the pullback on a random cotangent 
    Δ = randn(eltype(val2), size(val2))
@@ -98,7 +97,7 @@ for basis in test_bases
    println_slim(@test ∂1[3] == ∂2[3])
 
    # look at gradients with respect to the parameters 
-   _foo = p -> dot(Δ, LuxCore.apply(l, X, p, st)[1])
+   _foo = p -> dot(Δ, LuxCore.apply(basis, X, p, st)[1])
    g1 = Zygote.gradient(_foo, ps)[1] 
    if sizeof(ps) == 0 
       println_slim(@test (isnothing(g1) || isempty(g1)))
@@ -110,7 +109,3 @@ for basis in test_bases
 
 end 
 
-
-##
-
-# @info("Test Second-order derivatices with Lux")