Describe the bug
NDTF.__init__ picks each ensemble member's input dimension, depth and temperature with
np.random.randint / np.random.normal at construction time. None of those draws are stored in the
config or the save bundle, so load() builds a differently-shaped network and load_state_dict
raises. Every NDTF artifact is unloadable — reproduced independently for NDTFRegressor;
NDTFClassifier and InferenceModel.from_path fail the same way.
Every other architecture in the zoo round-trips exactly (MLP, ResNet, FTTransformer, TabTransformer,
TabM, NODE, AutoInt, TabulaRNN, SAINT, ENODE, MambaTab, Mambular, MambAttention, Tangos, Trompt, MLPLSS).
NDTF can never be reloaded: the architecture redraws tree shapes from the global RNG at construction, so load() shape-mismatches
Where: deeptab/architectures/ndtf.py (70-86)
NDTF picks each ensemble member's input dimension, depth and temperature with np.random.randint/np.random.normal at init time, and none of those draws are stored in the config or the save bundle, so load() builds a differently-shaped network and load_state_dict raises — every NDTF artifact is unloadable.
Observed: RuntimeError: Error(s) in loading state_dict for TaskModel:
size mismatch for estimator.trees.0.inner_nodes.0.weight: copying a param with shape torch.Size([4095, 47]) from checkpoint, the shape in current model is torch.Size([255, 47]).
size mismatch for estimator.trees.1.inner_nodes.0.weight: copying a param with shape torch.Size([2047, 2]) ... torch.Size([16383, 25]).
(~3.9 KB of size mismatches, one pair per tree). NDTFClassifier fails identically; InferenceModel.from_path on the same artifact fails too.
Expected: NDTFRegressor.load(path) returns a model that reproduces the saved predictions, like every other architecture in the sweep (MLP, ResNet, FTTransformer, TabTransformer, TabM, NODE, AutoInt, TabulaRNN, SAINT, ENODE, MambaTab, Mambular, MambAttention, Tangos, Trompt and MLPLSS all round-trip exactly).
Repro
import warnings; warnings.simplefilter('ignore')
import numpy as np, pandas as pd
from deeptab.models.ndtf import NDTFRegressor
from deeptab.configs import TrainerConfig
rng = np.random.default_rng(0)
X = pd.DataFrame({'a': rng.normal(size=48), 'b': rng.normal(size=48), 'c': rng.integers(0,5,48).astype(float)})
y = X['a'].values * 2
m = NDTFRegressor(trainer_config=TrainerConfig(max_epochs=1, batch_size=16, patience=2), random_state=0)
m.fit(X, y, accelerator='cpu', devices=1)
m.save('n.deeptab')
NDTFRegressor.load('n.deeptab') # RuntimeError
Expected behavior
NDTFRegressor.load(path) returns a model that reproduces the saved predictions. The per-tree draws
need to be persisted in the save bundle (or derived deterministically from the stored random_state)
so the reloaded architecture matches.
Screenshots
n/a
Desktop (please complete the following information):
- OS: macOS (Darwin 25.5.0, arm64)
- Python version: 3.11.15
- deeptab Version: 2.0.0 (main @ 4e6a359)
Additional context
torch 2.9.1, lightning 2.6.5, scikit-learn 1.9.0, numpy 2.4.6. Found in a second-pass review of v2.0.0
(seven independent lenses, each finding adversarially re-verified by a second reviewer, then re-run by
hand). Distinct from the already-filed #409-#426.
Describe the bug
NDTF.__init__picks each ensemble member's input dimension, depth and temperature withnp.random.randint/np.random.normalat construction time. None of those draws are stored in theconfig or the save bundle, so
load()builds a differently-shaped network andload_state_dictraises. Every NDTF artifact is unloadable — reproduced independently for
NDTFRegressor;NDTFClassifierandInferenceModel.from_pathfail the same way.Every other architecture in the zoo round-trips exactly (MLP, ResNet, FTTransformer, TabTransformer,
TabM, NODE, AutoInt, TabulaRNN, SAINT, ENODE, MambaTab, Mambular, MambAttention, Tangos, Trompt, MLPLSS).
NDTF can never be reloaded: the architecture redraws tree shapes from the global RNG at construction, so load() shape-mismatches
Where:
deeptab/architectures/ndtf.py(70-86)NDTF picks each ensemble member's input dimension, depth and temperature with
np.random.randint/np.random.normalat init time, and none of those draws are stored in the config or the save bundle, soload()builds a differently-shaped network andload_state_dictraises — every NDTF artifact is unloadable.Observed: RuntimeError: Error(s) in loading state_dict for TaskModel:
size mismatch for estimator.trees.0.inner_nodes.0.weight: copying a param with shape torch.Size([4095, 47]) from checkpoint, the shape in current model is torch.Size([255, 47]).
size mismatch for estimator.trees.1.inner_nodes.0.weight: copying a param with shape torch.Size([2047, 2]) ... torch.Size([16383, 25]).
(~3.9 KB of size mismatches, one pair per tree). NDTFClassifier fails identically; InferenceModel.from_path on the same artifact fails too.
Expected:
NDTFRegressor.load(path)returns a model that reproduces the saved predictions, like every other architecture in the sweep (MLP, ResNet, FTTransformer, TabTransformer, TabM, NODE, AutoInt, TabulaRNN, SAINT, ENODE, MambaTab, Mambular, MambAttention, Tangos, Trompt and MLPLSS all round-trip exactly).Repro
Expected behavior
NDTFRegressor.load(path)returns a model that reproduces the saved predictions. The per-tree drawsneed to be persisted in the save bundle (or derived deterministically from the stored
random_state)so the reloaded architecture matches.
Screenshots
n/a
Desktop (please complete the following information):
Additional context
torch 2.9.1, lightning 2.6.5, scikit-learn 1.9.0, numpy 2.4.6. Found in a second-pass review of v2.0.0
(seven independent lenses, each finding adversarially re-verified by a second reviewer, then re-run by
hand). Distinct from the already-filed #409-#426.