Skip to content

[BUG] NDTF models can never be reloaded — the architecture redraws its tree shapes from the global RNG #444

Description

@ChrisW09

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.

Metadata

Metadata

Assignees

No one assigned

    Labels

    bugSomething isn't working

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions