Add temporal embedding and ModernNCA model

[AdityaPandeyCN]

This PR adds temporal embedding feature and ModernNCA model. It also aims to improve docstrings following julia convention.

[AdityaPandeyCN]

@jeremiedb Initially I thought I had to refactor the doc-strings matching julia convention and PR to the jdb/mnca branch but that branch dont have the changes that were merged earlier moreover it fails to load too.

Are these cleanups from your side or just incomplete syncs ? , I am confused on this.
In this PR I have added the changes back, and also found a middle ground for the TabM coupling: instead of TabM's build_chain calling into the embedding module (per_feature_widths/has_real_embedding), I moved that method to the Models layer, which already imports both TabM and Embeddings. So module TabM no longer depends on the embedding module, while the grouped scaling-init still works (only for TabM, via dispatch).

[jeremiedb]

Sorry about the sync confusion, jdb/mnca branch ended-up with conflicts following an attempt to sync with the v4 branch that is now on main. Having your PR open against lux-v4 like here (or lux-new which has been sync with main).
lux-v4 and lux-new are the both proper base for your PR.

[AdityaPandeyCN]

@jeremiedb Please take a look when time permits.

[jeremiedb]

The Embedding layer still appear heavy at first sight.
For example, the embedding_config in the learner (NeuroTabRegressor), for a simple BatchNorm embedding under current/lux-new:
NeuroTabModels.Models.Embeddings.EmbeddingConfig(:batchnorm, 1, :relu, 32, 32, 0.01f0)
and under this PR:
NeuroTabModels.Models.Embeddings.EmbeddingLayer{NeuroTabModels.Models.Embeddings.BatchNormEmbeddings, Nothing}(NeuroTabModels.Models.Embeddings.BatchNormEmbeddings(), nothing)

At least, I think that the docs for the Embedding layer should provide a more explicit description of the design; how one can either provide an embeding config or a dict.
For the embeddings layers in general, I think that sticking to kwargs for the construction of each of the embedding layer's chain. For example, below for linear embed, why not have nfeats::Int, _x and nfeats, s.d_embedding as kwarg?:

NeuroTabModels.jl/src/models/embeddings/config.jl

Lines 291 to 292 in ebeda97

	_build_num(s::LinearEmbeddings, nfeats::Int, _x) =
	Chain(_LinearEmbeddings(nfeats, s.d_embedding; activation=act_dict[s.activation]), FlattenLayer())

Also, why not have _LinearEmbeddings build the complete Chain, so there's no need wrap the output of the _LinearEmbeddings() call into a chain that adds a FlattenLayer step?

Could the the dispatch for Nothing be dropped?

NeuroTabModels.jl/src/models/embeddings/config.jl

Line 238 in ebeda97

_num_from_dict(::Nothing, d) = nothing

NeuroTabModels.jl/src/models/embeddings/config.jl

Line 280 in ebeda97

needs_x_train(::Nothing) = false

For the Identity embedding, could it be simpler to use the utility Lux.NoOpLayer?: https://lux.csail.mit.edu/stable/api/Lux/layers#Misc.-Helper-Layers

I understand the current general design where there's no full independency between the embedding and associated architecture adds a layers of complexity. And that the not having full info from the learner (NeuroTabRegressor/Classifier) to build the chain (in cases like piecewise where actual data needs to be seen, or more geenrally for nfeats) also prevent a lightweight implementation. I think this PR made good progress in best accomodating these constraints. I'd still see a need to clear up the few elements raised above before being in a release state. Longer term; I'll look to revisit the design such as to make a convention for the chain to be fully derivable from from the learner.

[jeremiedb]

In relation to the previous general comment: could we make the EmbeddingLayer defaults to a IdentityEmbedding? So there cannot be a "Nothing" element for it. Then, if there's a temporal embedding (or eventual cat embedding as well), then the Embedding layer performs a vcat of these parallel / branch embedding layers: embed = vcat(numericalembedding(x), temporalembedding(x))

[jeremiedb]

Looks good!
With these latest embedding refactor, I think the PR is in good shape.
Just a little tweak for "testmode".
Also, please add an embedding page to the docs.

[jeremiedb]

add st = LuxCore.testmode(st), otherwise returns warning of running state in train mode without gradient tracking

[jeremiedb]

@AdityaPandeyCN could you look at these small remaining items prior to merging this PR and closing the NeuroTabModels story?

[AdityaPandeyCN]

Addressed them.