trainable AORlayer + some bug fixed

[CheukHinHoJerry]

This PR mainly fixes some bugs and implements the AORlayer for trainable $\zeta$ in Gaussian basis. The problem of matrix multiplication with PtrArray{Hyper} is resolved from Octavian.jl's end.

Remainder on we can actually do something like this:

function (l::LuxLayer{BasisType})(X, ps, st)
     # alloc B
     evaluate!(B, basis::BasisType, X, ps)
    return B
end

evaluate!(B, basis::BasisType, X) = evaluate!(B, basis::BasisType, X, basis.ps, basis.st)

function evaluate!(B, basis::BasisType, X, ps, st)
   # this is the main evaluation code.
end

Maybe we can leave this PR open and use the dev branch to implement our code in ACEpsi.jl until we have a nice clean up with the Lux interface?

CC @cortner @DexuanZhou

[cortner]

The states also need to get passed in. If not the states then the temporary array cache. Yes, maybe that's better actually.

[cortner]

What's the status of this PR?

[CheukHinHoJerry]

I think we talked about this today. I will try to do what we have discussed for trainable basis within this week. I am doing this pr together with https://github.com/ACEsuit/Polynomials4ML.jl/issues/45#issue-1765806905. Basically updates includes:

• [ ] migrating to evaluate!(P, basis, x, zeta, pool) interface and allow the use of tmparray in lux layers for non-allocating lux layer, which is to use tmp array in the middle layers and possible an FlexArrayCache in the last layer or we can do a copy.
• [x] trainable AORbasis with pullbacks
• [x] type instability fixes and extra tests added to make sure Zygote.gradient is calling correct overloaded pullback

For forwarddiff on zeta, I suggest not to include this here since it would be something that we have to test and play around with for a while to see how it actually goes?

[CheukHinHoJerry]

@cortner This is ready to be merged.

[cortner]

The only thing I don't really like about this PR is that it mixes a few unrelated issues - e.g. the non-allocating interface could have been treated in a separate PR. I don't think it is worth trying to separate them out now, that would be too much work. But in future, let's please try and make smaller PRs that solve discrete problems and are easier to review.

[CheukHinHoJerry]

Thank you for very review and comment. I will make sure that each PR only take care of a small and related issues in the future.

[cortner]

Does this PR need any new docs?

[cortner]

Is this a patch version?

[CheukHinHoJerry]

Re Docs: I think I need to write a documentation for the lux layer and interfaces on pullbacks. Maybe it is also nice to have a small hierarchy diagram to explain what is ScalarP4ML etc...

Should I do this right away?

Re patch version : Yes I think so.

[cortner]

I think I need to write a documentation for the lux layer and interfaces on pullbacks. Maybe it is also nice to have a small hierarchy diagram to explain what is ScalarP4ML etc...

At this point, I don't think this is user-facing yet. Let's do this once we finalize the pullback interfaces and make those implementations backward compatible.

[cortner]

I tested on 1.9 and 1.10 and will merge and register.

[cortner]

I've submitted this to General as version 0.1.5

[CheukHinHoJerry]

Thank you very much for your review.

[cortner]

this is now registered.