Refactoring and Cleanup of Interface

[cortner]

The purpose of this PR started as an experiment moving the entire P4ML package over to using Bumper and WithAlloc and removing all usage of ObjectPools. It now includes

• fully retiring OP in favour of Bumper
• cleanup of the pb and pfwd interface
• remove all hand-written pb2 and pwfd
• cleanup of the type hierarchy
• cleanup of allocating / non-allocating interface
• cleanup sphericart interface and retire the P4ML 3D harmonics implementations
• some improvements to unified testing

My sense now is that the extra cost of allocations in typical ACE models will cost a moderate % of runtime, somewhere between 30-50% typically. This is more than acceptable for prototyping. For fast final models we can then hand-write the chains to get the best possible performance.

I'm unifying and finalizing the pb/pfwd interface while I do this so I can write as much shared code as possible.

TODO

• decide on whether we support laplacians in P4ML or only the ability to use hyper-duals. It is not clear to me what is the right choice.
• bring the linear layer back (as discussed with Jerry)
• make sure the sparseproduct has the functionality needed elsewhere.
• cleanup and check the lux interface
• confirm the sphericart wrapper serves our purpose
• rename -> pullback! and pushforward!
• try and test unified second pullback via ForwardDiff
• try unified pushforward implementation via ForwardDiff
• allocation tests for most embeddings and all tensors (can bring in more over time)
• Extend withalloc to PooledSparseProduct

Leave for future PRs

• complex bases seem to fail allocation tests
• improve the unified testing functions (a lot of testing functionality is duplicated throughout the test_*** scripts
• provide literate examples how to compose manually and how to do pfwd and pb2 through such models, compare with Lux-style implementations of the same models, advantages and disadvantages
• literate example how to compute a laplacian
• update benchmark code

[CheukHinHoJerry]

Should we also remove anything related to FlexArray?

[cortner]

Yes for sure. Most of it should be gone. I just haven't reintegrated all parts of the codebase yet.

[CheukHinHoJerry]

SparseProduct looks fine elsewhere. The functionality now it provides is sufficient. I will have a final clean up now.

[cortner]

Thank you!

[cortner]

WithAlloc.jl is now registered in the general registry, which means we could start thinking about merging this PR and registering the next minor version? @CheukHinHoJerry , what is still missing?

[CheukHinHoJerry]

Sorry for late reply - I had some local changes on cleaning up the interface of SparseProduct that unifies with PooledSparseProduct. Just pushed them and now I think it is ready if the CI passes.

[CheukHinHoJerry]

CI stucks at testing SparseProduct, which probably relates to #87. That doesn't happen for me previously but it suddenly occurs for the last commit.

[cortner]

Ok, I'll take a look and see if I can fix it.

[CheukHinHoJerry]

I think I see why - I will fix it now.

[cortner]

Thank you, Jerry!

[cortner]

So maybe my last question for @DexuanZhou and @CheukHinHoJerry -- do you need laplacians HERE in P4ML or is it enough for you that you can use P4ML layers with Duals and HyperDuals?

[CheukHinHoJerry]

I think it is enough for us just to make evaluation of Duals and HyperDuals work. If there is a feature request or there are more than one projects that needs laplacian we can then consider adding that?

[cortner]

so we are actually ready to merge this? If you confirm, I'll go ahead with it :)

[CheukHinHoJerry]

Yes please I am happy for this to be merged.

[cortner]

Everyone - thank you for your feedback and help with this PR. I've now registereed this as v0.3.0 in the ACE registry (not in General yet since I'm having a problem with sphericart...)

[cortner]

This is now also registered in General