Embed and pool

[CheukHinHoJerry]

Implements PooledEmbeddings for double pullback / reverse over reverse

still WIP

Some notes for myself, we assume that PooledEmbedding has the following struct:

struct PooledEmbedding{NB, TB <: Tuple} <: AbstractPoly4MLBasis
   embeddings::TB
   pooling::PooledSparseProduct{NB}
   @reqfields()
end

and then TB is a tuple of embeddings with

evaluate(embedding, X)

and

ChainRulesCore.rrule(evaluate, embedding, X)

defined.

With X isa State and X::State + Y::State also defined (possibly for positions? I don't understand how operations like+ is defined in State yet so I will have to think over this but it should be straight forward.).

[CheukHinHoJerry]

Some how this is working, but there are some issues on type constrains/multiple dispatch that I have to mull over. fdtest coming soon tmr.

[CheukHinHoJerry]

A note to myself why we are usnig ChainRulesCore in:

      @nexprs $NB i -> begin
         _, embed_pb_i = ChainRulesCore.rrule(evaluate, basis.embeddings[i], X)    
      end
      # pooling_pb : Vec -> (B_1, B_2, ..., B_NB)
      # embed_pb_i : B_i -> ∂X
      # use 3 since the interface must return (NoTangent(), NoTangent(), ∂)
      ∂X = zero(X)
      ∂BBs = pooling_pb(Δ)[3]
      # writes and accumulate to ∂X
      @nexprs $NB i -> begin
         ∂X_i = embed_pb_i(∂BBs[i])
         ∂X .+= ∂X_i[3]
      end

Since Zygote will not be a dependency of P4ML.

[CheukHinHoJerry]

somehow the reverse over reverse is working on EquivariantModels.jl's end. But then something like this

      @nexprs $(NB) i -> begin
         ∂X_i = embed_pb_i(∂BBs[i])
         ∂X .+= ∂X_i[3]
      end

is not general if .+ is not defined on X. But from my end I prefer to clean up and then merge and fix it with another PR after the State package is available.

[CheukHinHoJerry]

@cortner I think this is ready. There are some generalization to the State idea but I think it generally goes to overriding functions properly independent of the implementation of the current PR.

[cortner]

this seems no longer needed, so we are closing it, but keep the branch in case we ever want to revive it -- e.g. it could be a useful performance optimisation at some point.