Check input dimensions

python/nvforest/nvforest/_base.py

@@ -90,6 +90,11 @@ def apply(
         """
         pass
 
+    @property
+    @abstractmethod
+    def num_features(self) -> int:
+        pass
+
     @property
     @abstractmethod
     def num_outputs(self) -> int:

python/nvforest/nvforest/_forest_inference.py

@@ -387,6 +387,10 @@ def apply(
     ) -> DataType:
         return self.forest.apply(X, chunk_size=chunk_size)
 
+    @property
+    def num_features(self) -> int:
+        return self.forest.num_features
+
     @property
     def num_outputs(self) -> int:
         return self.forest.num_outputs
@@ -488,6 +492,10 @@ def apply(
     ) -> DataType:
         return self.forest.apply(X, chunk_size=chunk_size)
 
+    @property
+    def num_features(self) -> int:
+        return self.forest.num_features
+
     @property
     def num_outputs(self) -> int:
         return self.forest.num_outputs
@@ -603,6 +611,10 @@ def apply(
     ) -> DataType:
         return self.forest.apply(X, chunk_size=chunk_size)
 
+    @property
+    def num_features(self) -> int:
+        return self.forest.num_features
+
     @property
     def num_outputs(self) -> int:
         return self.forest.num_outputs
@@ -706,6 +718,10 @@ def apply(
     ) -> DataType:
         return self.forest.apply(X, chunk_size=chunk_size)
 
+    @property
+    def num_features(self) -> int:
+        return self.forest.num_features
+
     @property
     def num_outputs(self) -> int:
         return self.forest.num_outputs

python/nvforest/nvforest/detail/forest_inference.pyx

@@ -355,12 +355,22 @@ class ForestInferenceImpl:
     def elem_postprocessing(self) -> str:
         return self.impl.elem_postprocessing()
 
+    def _validate_input_dims(self, X: DataType) -> None:
+        if len(X.shape) != 2:
+            raise ValueError("Expected a 2D array for X")
+        if X.shape[1] != self.num_features:
+            raise ValueError(
+                f"Expected {self.num_features} features in the input "
+                f"but X has {X.shape[1]} features"
+            )
+
     def predict(
         self,
         X: DataType,
         *,
         chunk_size: Optional[int] = None,
     ) -> DataType:
+        self._validate_input_dims(X)
         # Returns probabilities if the model is a classifier
         return self.impl.predict(
             X, chunk_size=(chunk_size or self.default_chunk_size)
@@ -372,6 +382,7 @@ class ForestInferenceImpl:
         *,
         chunk_size: Optional[int] = None,
     ) -> DataType:
+        self._validate_input_dims(X)
         chunk_size = (chunk_size or self.default_chunk_size)
         return self.impl.predict(
             X, predict_type="per_tree", chunk_size=chunk_size
@@ -383,6 +394,7 @@ class ForestInferenceImpl:
         *,
         chunk_size: Optional[int] = None,
     ) -> DataType:
+        self._validate_input_dims(X)
         chunk_size = (chunk_size or self.default_chunk_size)
         return self.impl.predict(
             X, predict_type="leaf_id", chunk_size=chunk_size

python/nvforest/tests/test_nvforest.py

@@ -856,3 +856,29 @@ def test_wide_data():
     # Inference should run without crashing
     fm = nvforest.load_from_sklearn(clf)
     _ = fm.predict(X)
+
+
+@pytest.mark.parametrize("input_size", [4, 6], ids=["too_narrow", "too_wide"])
+@pytest.mark.parametrize(
+    "predict_func",
+    [
+        nvforest.CPUForestInferenceClassifier.predict,
+        nvforest.CPUForestInferenceClassifier.predict_per_tree,
+        nvforest.CPUForestInferenceClassifier.apply,
+    ],
+    ids=["predict", "predict_per_tree", "apply"],
+)
+def test_incorrect_data_shape(input_size, predict_func):
+    n_rows = 50
+    n_features = 5
+    X = np.random.normal(size=(n_rows, n_features)).astype(np.float32)
+    y = np.asarray([0, 1] * (n_rows // 2), dtype=np.int32)
+
+    clf = RandomForestClassifier(max_features="sqrt", n_estimators=10)
+    clf.fit(X, y)
+
+    fm = nvforest.load_from_sklearn(clf, device="cpu")
+    assert fm.num_features == n_features
+    with pytest.raises(ValueError, match=f"Expected {n_features} features"):
+        X_test = np.zeros((1, input_size))
+        _ = predict_func(fm, X_test)