fix(architectures): build NDTF on a one-dimensional input

deeptab/architectures/ndtf.py

@@ -84,7 +84,8 @@ def __init__(
         if input_dimensions is None:
             input_dimensions = [input_dim]
             for _ in range(self.hparams.n_ensembles - 1):
-                input_dimensions.append(np.random.randint(1, input_dim))
+                # A one-dimensional input has no smaller prefix to sample.
+                input_dimensions.append(np.random.randint(1, input_dim) if input_dim > 1 else 1)
         if depths is None:
             depths = [
                 np.random.randint(self.hparams.min_depth, self.hparams.max_depth)

tests/architectures/test_ndtf.py

@@ -45,6 +45,13 @@ def test_feature_projection_matches_convolution_outputs_and_gradients(width, bat
     torch.testing.assert_close(model.conv_layer.bias.grad, convolution.bias.grad)
 
 
+def test_one_dimensional_input_builds_every_tree_on_that_input():
+    num_info = {"feature": {"preprocessing": "none", "dimension": 1, "categories": None}}
+    model = NDTF(feature_information=(num_info, {}, {}), num_classes=1, config=NDTFConfig(n_ensembles=3))
+    assert model.input_dimensions == [1, 1, 1]
+    assert model([torch.randn(5, 1)], [], []).shape == (5, 1)
+
+
 @pytest.mark.parametrize("device", ["cpu", "mps"])
 @pytest.mark.parametrize("with_penalty", [False, True])
 def test_forest_backward_has_finite_gradients(device, with_penalty):