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

• [x] 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.
• [x] bring the linear layer back (as discussed with Jerry)
• [x] make sure the sparseproduct has the functionality needed elsewhere.
• [x] cleanup and check the lux interface
• [x] confirm the sphericart wrapper serves our purpose
• [x] rename -> pullback! and pushforward!
• [x] try and test ~unified~ second pullback via ForwardDiff
• [x] try ~unified~ pushforward implementation via ForwardDiff
• [x] allocation tests for most embeddings and all tensors (can bring in more over time)
• [x] 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.

[cortner]

@CheukHinHoJerry -- in the spirit of slimming down the package, how do we feel about making the P4ML layers actual Lux layers rather than wrapping them into Lux layers? No rush to decide, but let's chat?

[CheukHinHoJerry]

I quite like the idea of having P4ML.lux since I can then define a layer wrapper for my basis <: an abstract layer type in P4ML and then use all the functionalities conveniently even in other packages in a clean way so I prefer keeping that on the first thought. But let me mull over it and we can discuss tmr? I plan to come to campus.

[cortner]

Sounds good to me. It might actually make it easier to ensure that functinality doesn't start interfering - and we might be able to re-implement the "release" functionality.

[cortner]

I think I've now done everything I want to do for now. I propose the following:

• @CheukHinHoJerry and @DexuanZhou , can you please take a look at the TODO lists at the top of the PR and comment if you have any concerns
• once we agreed we have a good candiate for the next minor version, I branch main -> v0.2.x-dev and then merge this PR into main i.e. it becomes the new development branch. I then register it as 0.3.0. Any bugs, problems etc we find we can fix over the course of 0.3.x development.

[CheukHinHoJerry]

Looks good. I don't have any concern with the suggestions either.

[cortner]

@CheukHinHoJerry -- thank you. Would you mind then taking a look at the linear layer and SparseProduct?

[cortner]

And maybe we can work together on a model building demo.

[CheukHinHoJerry]

Yes I will take a look at the linear layer first and then the sparseproduct over the weekend.

[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