Resolve most of the issues in comments

ACEsuit · Oct 5, 2023 · b363596 · b363596
1 parent a115a63
commit b363596
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 7 deletions.
diff --git a/src/ConstLinearLayer.jl b/src/ConstLinearLayer.jl
@@ -2,17 +2,15 @@ import ChainRulesCore: rrule
 using LuxCore
 using LuxCore: AbstractExplicitLayer
 
-struct ConstLinearLayer{T} <: AbstractExplicitLayer # where {in_dim,out_dim,T}
-    W::AbstractMatrix{T}  
+struct ConstLinearLayer <: AbstractExplicitLayer # where {in_dim,out_dim,T}
+    W #::AbstractMatrix{T}  
     position::Union{Vector{Int64}, UnitRange{Int64}}
-    in_dim::Integer
-    out_dim::Integer
 end
 
-ConstLinearLayer(W::AbstractMatrix{T}) where T = ConstLinearLayer(W,1:size(W,2),size(W,2),size(W,1))
-ConstLinearLayer(W::AbstractMatrix{T}, pos::Union{Vector{Int64}, UnitRange{Int64}}) where T = ConstLinearLayer(W,pos,size(W,2),size(W,1))
+ConstLinearLayer(W) where T = ConstLinearLayer(W,1:size(W,2))
+# ConstLinearLayer(W, pos::Union{Vector{Int64}, UnitRange{Int64}}) = ConstLinearLayer(W,pos)
 
-(l::ConstLinearLayer)(x::AbstractVector) = l.in_dim == length(x[l.position]) ? l.W * x[l.position] : error("x (or the position index) has a wrong length!")
+(l::ConstLinearLayer)(x::AbstractVector) = l.W * x[l.position]
 
 (l::ConstLinearLayer)(x::AbstractMatrix) = begin
     Tmp = l(x[1,:])

diff --git a/src/builder.jl b/src/builder.jl
@@ -261,6 +261,7 @@ function equivariant_model(spec_nlm, radial::Radial_basis, L::Int64; categories=
       C = _rpi_A2B_matrix(cgen, spec_nlm)
    end
 
+   #TODO:make use [ C[i], pos[i] ] to generate another sparse matrix so that...
    l_sym = islong ? Lux.Parallel(nothing, [ConstLinearLayer(C[i],pos[i]) for i = 1:L+1]... ) : ConstLinearLayer(C)
    # C - A2Bmap
    luxchain = append_layer(luxchain_tmp, l_sym; l_name = :BB)