Lux new

benchmarks/titanic-logloss.jl

@@ -85,6 +85,12 @@ learner = NeuroTabRegressor(
     device=:cpu,
     backend=:reactant
 )
+learner.embedding_config
+# main / lux-new
+# NeuroTabModels.Models.Embeddings.EmbeddingConfig(:batchnorm, 1, :relu, 32, 32, 0.01f0)
+# mnca/temporal
+# NeuroTabModels.Models.Embeddings.EmbeddingLayer{NeuroTabModels.Models.Embeddings.BatchNormEmbeddings, Nothing}(NeuroTabModels.Models.Embeddings.BatchNormEmbeddings(), nothing)
+# NeuroTabModels.Models.Embeddings.EmbeddingLayer{NeuroTabModels.Models.Embeddings.PiecewiseLinearEmbeddings, Nothing}(NeuroTabModels.Models.Embeddings.PiecewiseLinearEmbeddings(8, 16, :identity, :B), nothing)
 
 @time m = NeuroTabModels.fit(
     learner,

docs/make.jl

@@ -7,6 +7,7 @@ pages = [
     "Quick start" => "quick-start.md",
     "API" => "API.md",
     "Design" => "design.md",
+    "Embeddings" => "embeddings.md",
     "Models" => [
         "Interface" => "models/models.md",
         "MLP" => "models/mlp.md",

docs/src/embeddings.md

@@ -1,5 +1,43 @@
 # Embeddings
 
+Embeddings transform raw input columns before they reach a model backbone. Numerical
+embeddings operate column-wise, while temporal embeddings reserve one column for
+time features and concatenate that branch with the remaining numerical features.
+
+## Available Embeddings
+
+- `IdentityEmbedding`: leaves numerical features unchanged.
+- `LinearEmbeddings`: projects each numerical feature to a learned vector.
+- `PeriodicEmbeddings`: expands each feature with learned sinusoidal terms before projection.
+- `PiecewiseLinearEmbeddings`: computes feature bins from the training data and embeds the resulting piecewise-linear encoding.
+- `BatchNormEmbeddings`: batch-normalizes raw numerical features without expanding them.
+- `TemporalEmbeddings`: embeds one time column with Fourier features and an optional trend term.
+
+## Configuration
+
+Use `EmbeddingLayer` to combine a numerical embedding with an optional temporal
+branch. A dictionary form is also available for model configuration dictionaries.
+
+```julia
+EmbeddingLayer(PeriodicEmbeddings(; d_embedding=24))
+
+EmbeddingLayer(;
+    num=PiecewiseLinearEmbeddings(; bins=32),
+    temp=TemporalEmbeddings(; index=1),
+)
+
+EmbeddingLayer(Dict(
+    :embedding_type => :periodic,
+    :d_embedding => 24,
+    :temporal => Dict(:index => 1),
+))
+```
+
+Piecewise-linear and temporal embeddings need training data when the embedding chain is
+built. Use `needs_x_train` to check that requirement for a spec.
+
+## API
+
 ```@autodocs
 Modules = [NeuroTabModels.Models.Embeddings]
 ```

ext/NeuroTabModelsReactantExt.jl

@@ -16,12 +16,17 @@ function _get_device(::Val{:reactant}, ::Val{D}; gpuID::Integer=0) where {D}
     return reactant_device()
 end
 
+_same_shape(a::AbstractArray, b::AbstractArray) = size(a) == size(b)
+_same_shape(a::Tuple, b::Tuple) =
+    length(a) == length(b) && all(_same_shape(ai, bi) for (ai, bi) in zip(a, b))
+_same_shape(_, _) = false
+
 function _infer_loop(::Val{:reactant}, chain, data, x0, dev, cdev, ps, st)
     compiled = @compile _forward_reduce(chain, dev(x0), ps, st)
 
     preds = Vector{AbstractArray}()
     for x in data
-        if size(x) == size(x0)
+        if _same_shape(x, x0)
             pred = compiled(chain, dev(x), ps, st)
         else
             pred = Reactant.@jit _forward_reduce(chain, dev(x), ps, st)

src/Fit/callback.jl

@@ -7,6 +7,7 @@ using ..Learners: LearnerTypes
 using ...Infer: reduce_pred, _get_device
 using ..Data: get_df_loader_train
 using ..Metrics
+using ...Models
 
 using Lux: Training, testmode
 
@@ -26,7 +27,8 @@ end
 function CallBack(
     config::LearnerTypes,
     df::AbstractDataFrame,
-    cache;
+    cache,
+    m;
     feature_names,
     target_name,
     weight_name=nothing,
@@ -43,6 +45,7 @@ function CallBack(
     deval = get_df_loader_train(dfg; feature_names, target_name, weight_name, offset_name, scalers, batchsize, shuffle=false) |> dev
 
     ps, st = ts.parameters, testmode(ts.states)
+    deval = Models.eval_dataloader(m.chain, m.info, deval, dev, ps, st)
     d0 = first(deval)
     eval_compiled = _build_eval_step(ts.model, feval, d0, ps, st; reactant=config.backend == :reactant)
 

src/Fit/fit.jl

@@ -9,7 +9,7 @@ using ..Losses
 using ..Metrics
 using ..Infer: reduce_pred, _get_device
 
-import Random: Xoshiro
+import Random: Xoshiro, default_rng
 import Statistics: mean, std
 import MLJModelInterface: fit
 import Optimisers: OptimiserChain, WeightDecay, NAdam, Adam
@@ -73,19 +73,12 @@ function init(
 
     # Build chain: optional embeddings + architecture backbone
     embed_config = config.embedding_config
-    if isnothing(embed_config)
-        chain = config.arch(; nfeats, outsize)
-    else
-        if embed_config.embedding_type == :piecewise
-            x_train = Matrix{Float32}(df[:, feature_names])
-        else
-            x_train = nothing
-        end
-        embed_chain = embed_config(; nfeats, x_train)
-        d_in = nfeats * embed_config.d_embedding
-        d_features = fill(embed_config.d_embedding, nfeats)
-        chain = Chain(embed_chain, config.arch(; nfeats=d_in, outsize, d_features, scaling_init_override=:normal))
-    end
+    x_train = Models.Embeddings.needs_x_train(embed_config) ? Matrix{Float32}(df[:, feature_names]) : nothing
+    embed_chain = Models.Embeddings.build_embedding_chain(embed_config, nfeats; x_train)
+    d_in = Models.Embeddings.embedding_width(embed_chain, randn(Float32, nfeats, 2), default_rng())
+    chain = Models.build_chain(config.arch, embed_chain;
+        nfeats, outsize, d_in, embedding_config=embed_config, loss_type=L)
+
 
     info = Dict(
         :nrounds => 0,
@@ -105,6 +98,8 @@ function init(
 
     rng = Xoshiro(config.seed)
     ps, st = Lux.setup(rng, m.chain) |> dev
+    data = Models.train_dataloader(config.arch, m, data, df;
+        feature_names, target_name, loss_type=L, scalers, batchsize, dev, rng)
     opt = OptimiserChain(NAdam(config.lr), WeightDecay(config.wd))
     ts = Training.TrainState(m.chain, ps, st, opt)
 
@@ -164,7 +159,7 @@ function fit(
 
     logger = nothing
     if !isnothing(deval)
-        cb = CallBack(config, deval, cache; feature_names, target_name, weight_name, offset_name, group_key)
+        cb = CallBack(config, deval, cache, m; feature_names, target_name, weight_name, offset_name, group_key)
         logger = init_logger(config)
         cb(logger, 0, cache[:train_state])
         (verbosity > 0) && @info "Init training" metric = logger[:metrics][end]

src/learners.jl

@@ -5,14 +5,19 @@ import MLJModelInterface: fit, update, predict, schema
 import Random
 
 using ..Models
-using ..Models: EmbeddingConfig
+using ..Models: AbstractEmbedding, IdentityEmbedding, EmbeddingLayer
 export NeuroTabRegressor, NeuroTabClassifier, LearnerTypes
 
+_to_embedding(::Nothing)            = IdentityEmbedding()
+_to_embedding(e::AbstractEmbedding) = e
+_to_embedding(d::AbstractDict)      = EmbeddingLayer(d)
+_to_embedding(x) = error("`embedding_config` must be `nothing`, an `AbstractDict`, or an `AbstractEmbedding`; got $(typeof(x)).")
+
 mutable struct NeuroTabRegressor <: MMI.Deterministic
   loss::Symbol
   metric::Symbol
   arch::Architecture
-  embedding_config::Union{Nothing,EmbeddingConfig}
+  embedding_config::AbstractEmbedding
   nrounds::Int
   early_stopping_rounds::Int
   lr::Float32
@@ -47,8 +52,10 @@ A model type for constructing a NeuroTabRegressor, based on [NeuroTabModels.jl](
 - `backend=:zygote`:        Backend used by Lux. One of `:enzyme`, `:zygote`, or `:reactant`.
 - `device=:gpu`:            Execution device. One of `:cpu` or `:gpu`.
 - `gpuID=0`:                GPU device to use, only relevant if `device = :gpu`. `0` auto-selects.
-- `embedding_config=nothing`: Optional `Dict` or `EmbeddingConfig` for numerical feature embeddings.
-  E.g. `embedding_config=Dict(:embedding_type => :periodic, :d_embedding => 24)`.
+- `embedding_config=nothing`: Optional numerical/temporal embeddings. Accepts `nothing` (no-op),
+  an `AbstractEmbedding` (e.g. `EmbeddingLayer(num=PeriodicEmbeddings(d_embedding=24))`), or an
+  `AbstractDict` selecting the type via `:embedding_type` (e.g.
+  `Dict(:embedding_type => :periodic, :d_embedding => 24)`).
 
 `backend=:zygote` works on `:cpu` and `:gpu`; `backend=:reactant` works on `:cpu` and `:gpu` and uses Enzyme for AD.
 `backend=:enzyme` with `device=:gpu` is currently known to fail for some NeuroTabModels models.
@@ -194,11 +201,7 @@ function NeuroTabRegressor(arch::Architecture; kwargs...)
     @warn "`backend=:enzyme` with `device=:gpu` is currently known to fail for some NeuroTabModels models. Prefer `backend=:zygote` on GPU, `backend=:reactant` for Reactant, or `backend=:enzyme` on CPU."
   end
 
-  # Build EmbeddingConfig from Dict if needed
-  embed = args[:embedding_config]
-  if embed isa AbstractDict
-    embed = EmbeddingConfig(; embed...)
-  end
+  embed = _to_embedding(args[:embedding_config])
 
   config = NeuroTabRegressor(
     loss,
@@ -231,7 +234,7 @@ mutable struct NeuroTabClassifier <: MMI.Probabilistic
   loss::Symbol
   metric::Symbol
   arch::Architecture
-  embedding_config::Union{Nothing,EmbeddingConfig}
+  embedding_config::AbstractEmbedding
   nrounds::Int
   early_stopping_rounds::Int
   lr::Float32
@@ -259,7 +262,8 @@ A model type for constructing a NeuroTabClassifier, based on [NeuroTabModels.jl]
 - `backend=:zygote`:        Backend used by Lux. One of `:enzyme`, `:zygote`, or `:reactant`.
 - `device=:gpu`:            Execution device. One of `:cpu` or `:gpu`.
 - `gpuID=0`:                GPU device to use, only relevant if `device = :gpu`. `0` auto-selects.
-- `embedding_config=nothing`: Optional `Dict` or `EmbeddingConfig` for numerical feature embeddings.
+- `embedding_config=nothing`: Optional numerical/temporal embeddings. Accepts `nothing` (no-op),
+  an `AbstractEmbedding`, or an `AbstractDict` selecting the type via `:embedding_type`.
 
 `backend=:zygote` works on `:cpu` and `:gpu`; `backend=:reactant` works on `:cpu` and `:gpu` and uses Enzyme for AD.
 `backend=:enzyme` with `device=:gpu` is currently known to fail for some NeuroTabModels models.
@@ -390,11 +394,7 @@ function NeuroTabClassifier(arch::Architecture; kwargs...)
     @warn "`backend=:enzyme` with `device=:gpu` is currently known to fail for some NeuroTabModels models. Prefer `backend=:zygote` on GPU, `backend=:reactant` for Reactant, or `backend=:enzyme` on CPU."
   end
 
-  # Build EmbeddingConfig from Dict if needed
-  embed = args[:embedding_config]
-  if embed isa AbstractDict
-    embed = EmbeddingConfig(; embed...)
-  end
+  embed = _to_embedding(args[:embedding_config])
 
   config = NeuroTabClassifier(
     :mlogloss,