Add cuml.accel support for sklearn.ensemble.IsolationForest

cpp/include/cuml/ensemble/isolation_forest.hpp

@@ -145,9 +145,12 @@ void fit(const raft::handle_t& handle,
  * @param[in]  n_rows          Number of training samples
  * @param[in]  n_cols          Number of features
  * @param[in]  params          Hyperparameters (n_estimators, max_samples, max_depth, seed)
- * @param[out] c_normalization Normalization constant c(n) for the trained forest, needed to
- *                             turn average path lengths into anomaly scores
- * @param[in]  verbosity       Logging level
+ * @param[out] c_normalization    Normalization constant c(n) for the trained forest, needed to
+ *                                turn average path lengths into anomaly scores
+ * @param[out] feature_indices    Host buffer receiving each tree's sampled feature indices in
+ *                                row-major [n_estimators, resolved max_features] order
+ * @param[in]  feature_indices_size Number of elements available in feature_indices
+ * @param[in]  verbosity          Logging level
  */
 template <typename T>
 void fit_treelite(const raft::handle_t& handle,
@@ -157,6 +160,8 @@ void fit_treelite(const raft::handle_t& handle,
                   int n_cols,
                   const IF_params& params,
                   double* c_normalization,
+                  int* feature_indices,
+                  size_t feature_indices_size,
                   rapids_logger::level_enum verbosity = rapids_logger::level_enum::info);
 
 /**

cpp/src/isolation_forest/isolation_forest.cu

@@ -5,10 +5,12 @@
 
 #include "isolation_forest.cuh"
 
+#include <cuml/common/checked_arithmetic.hpp>
 #include <cuml/ensemble/isolation_forest.hpp>
 
 #include <raft/core/error.hpp>
 #include <raft/core/handle.hpp>
+#include <raft/util/cudart_utils.hpp>
 
 #include <rmm/device_uvector.hpp>
 #include <rmm/exec_policy.hpp>
@@ -161,13 +163,45 @@ void fit_treelite(const raft::handle_t& handle,
                   int n_cols,
                   const IF_params& params,
                   double* c_normalization,
+                  int* feature_indices,
+                  size_t feature_indices_size,
                   rapids_logger::level_enum verbosity)
 {
   ASSERT(c_normalization != nullptr, "Normalization output pointer cannot be null.");
 
   IsolationForestModel<T> forest;
   fit(handle, &forest, input, n_rows, n_cols, params, verbosity);
   *c_normalization = forest.c_normalization;
+
+  size_t expected_feature_indices =
+    ML::checked_mul<std::size_t>(forest.params.n_estimators, forest.n_features_per_tree);
+  ASSERT(feature_indices != nullptr || expected_feature_indices == 0,
+         "Feature indices output buffer cannot be null.");
+  ASSERT(feature_indices_size == expected_feature_indices,
+         "Expected feature indices output buffer of size %zu, got %zu.",
+         expected_feature_indices,
+         feature_indices_size);
+
+  if (forest.global_feature_indices.size() == 0) {
+    for (int tree = 0; tree < forest.params.n_estimators; ++tree) {
+      // Bounded by expected_feature_indices (validated above), so the per-row
+      // base offset is computed once rather than checked on every write.
+      size_t row_offset = static_cast<size_t>(tree) * forest.n_features_per_tree;
+      for (int feature = 0; feature < forest.n_features_per_tree; ++feature) {
+        feature_indices[row_offset + feature] = feature;
+      }
+    }
+  } else {
+    auto stream = handle.get_stream();
+    RAFT_CUDA_TRY(
+      cudaMemcpyAsync(feature_indices,
+                      forest.global_feature_indices.data(),
+                      ML::checked_mul<std::size_t>(expected_feature_indices, sizeof(int)),
+                      cudaMemcpyDeviceToHost,
+                      stream));
+    handle.sync_stream(stream);
+  }
+
   build_treelite_isolation_forest<T>(model_handle, handle, &forest);
 }
 
@@ -273,6 +307,8 @@ template CUML_EXPORT void fit_treelite<float>(const raft::handle_t&,
                                               int,
                                               const IF_params&,
                                               double*,
+                                              int*,
+                                              size_t,
                                               rapids_logger::level_enum);
 template CUML_EXPORT void fit_treelite<double>(const raft::handle_t&,
                                                TreeliteModelHandle*,
@@ -281,6 +317,8 @@ template CUML_EXPORT void fit_treelite<double>(const raft::handle_t&,
                                                int,
                                                const IF_params&,
                                                double*,
+                                               int*,
+                                               size_t,
                                                rapids_logger::level_enum);
 
 }  // namespace ML

cpp/tests/sg/isolation_forest_test.cu

@@ -482,6 +482,128 @@ TEST_F(IsolationForestTest, MaxFeaturesFitStoresOriginalFeatureIds)
   EXPECT_EQ(model.global_feature_indices.size(), 0);
 }
 
+TEST_F(IsolationForestTest, FitTreeliteReturnsSampledFeatureIds)
+{
+  const int n_samples    = 128;
+  const int n_features   = 8;
+  const int n_estimators = 5;
+  const int max_features = 3;
+
+  thrust::device_vector<float> X_rowmajor(n_samples * n_features);
+  thrust::device_vector<float> X_colmajor(n_samples * n_features);
+  raft::random::Rng rng(42);
+  rng.normal(X_rowmajor.data().get(), X_rowmajor.size(), 0.0f, 1.0f, stream);
+  handle->sync_stream(stream);
+  transpose_data(X_rowmajor, X_colmajor, n_samples, n_features);
+
+  IF_params params;
+  params.n_estimators = n_estimators;
+  params.max_samples  = 64;
+  params.max_features = max_features;
+  params.seed         = 42;
+
+  std::vector<int> feature_indices(static_cast<size_t>(n_estimators) * max_features);
+  TreeliteModelHandle tl_handle = nullptr;
+  double c_normalization        = 0.0;
+  fit_treelite(*handle,
+               &tl_handle,
+               X_colmajor.data().get(),
+               n_samples,
+               n_features,
+               params,
+               &c_normalization,
+               feature_indices.data(),
+               feature_indices.size());
+  ASSERT_NE(tl_handle, nullptr);
+  delete static_cast<tl::Model*>(tl_handle);
+
+  for (int tree = 0; tree < n_estimators; ++tree) {
+    auto first = feature_indices.begin() + static_cast<size_t>(tree) * max_features;
+    auto last  = first + max_features;
+    for (auto it = first; it != last; ++it) {
+      EXPECT_GE(*it, 0);
+      EXPECT_LT(*it, n_features);
+    }
+    std::vector<int> sorted(first, last);
+    std::sort(sorted.begin(), sorted.end());
+    EXPECT_EQ(std::unique(sorted.begin(), sorted.end()), sorted.end());
+  }
+}
+
+TEST_F(IsolationForestTest, FitTreeliteValidatesFeatureBufferSize)
+{
+  const int n_samples  = 32;
+  const int n_features = 4;
+
+  thrust::device_vector<float> X_rowmajor(n_samples * n_features);
+  thrust::device_vector<float> X_colmajor(n_samples * n_features);
+  raft::random::Rng rng(42);
+  rng.normal(X_rowmajor.data().get(), X_rowmajor.size(), 0.0f, 1.0f, stream);
+  handle->sync_stream(stream);
+  transpose_data(X_rowmajor, X_colmajor, n_samples, n_features);
+
+  IF_params params;
+  params.n_estimators = 2;
+  params.max_samples  = 16;
+  params.max_features = 2;
+
+  std::vector<int> feature_indices(3);
+  TreeliteModelHandle tl_handle = nullptr;
+  double c_normalization        = 0.0;
+  EXPECT_THROW(fit_treelite(*handle,
+                            &tl_handle,
+                            X_colmajor.data().get(),
+                            n_samples,
+                            n_features,
+                            params,
+                            &c_normalization,
+                            feature_indices.data(),
+                            feature_indices.size()),
+               raft::exception);
+  EXPECT_EQ(tl_handle, nullptr);
+}
+
+TEST_F(IsolationForestTest, FitTreeliteReturnsFullFeatureRange)
+{
+  const int n_samples    = 32;
+  const int n_features   = 4;
+  const int n_estimators = 2;
+
+  thrust::device_vector<float> X_rowmajor(n_samples * n_features);
+  thrust::device_vector<float> X_colmajor(n_samples * n_features);
+  raft::random::Rng rng(42);
+  rng.normal(X_rowmajor.data().get(), X_rowmajor.size(), 0.0f, 1.0f, stream);
+  handle->sync_stream(stream);
+  transpose_data(X_rowmajor, X_colmajor, n_samples, n_features);
+
+  IF_params params;
+  params.n_estimators = n_estimators;
+  params.max_samples  = 16;
+  params.max_features = n_features;
+  params.seed         = 42;
+
+  std::vector<int> feature_indices(static_cast<size_t>(n_estimators) * n_features);
+  TreeliteModelHandle tl_handle = nullptr;
+  double c_normalization        = 0.0;
+  fit_treelite(*handle,
+               &tl_handle,
+               X_colmajor.data().get(),
+               n_samples,
+               n_features,
+               params,
+               &c_normalization,
+               feature_indices.data(),
+               feature_indices.size());
+  ASSERT_NE(tl_handle, nullptr);
+  delete static_cast<tl::Model*>(tl_handle);
+
+  for (int tree = 0; tree < n_estimators; ++tree) {
+    for (int feature = 0; feature < n_features; ++feature) {
+      EXPECT_EQ(feature_indices[static_cast<size_t>(tree) * n_features + feature], feature);
+    }
+  }
+}
+
 TEST_F(IsolationForestTest, ConstantFeaturesDoNotStopSplitting)
 {
   const int n_samples    = 64;

docs/source/cuml-accel/compatibility.rst

@@ -225,6 +225,22 @@ To compare results between estimators, we recommend comparing scores like
    - If ``y`` is a multi-output target.
 
 
+.. dropdown:: ``IsolationForest``
+   :name: isolationforest
+
+   ``IsolationForest`` will fall back to CPU in the following cases:
+
+   - If ``warm_start=True``.
+   - If a non-``None`` ``sample_weight`` is passed to ``fit`` or
+     ``fit_predict``.
+   - If ``X`` is sparse.
+   - If ``X`` contains missing or non-finite values.
+
+   Additionally, the following fitted attributes are currently not computed:
+
+   - ``estimators_samples_``
+
+
 sklearn.kernel_ridge
 ~~~~~~~~~~~~~~~~~~~~
 

python/cuml/cuml/accel/_overrides/sklearn/ensemble.py

@@ -1,14 +1,20 @@
 #
-# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 
+import numpy as np
+
 import cuml.ensemble
 from cuml.accel.estimator_proxy import ProxyBase
 from cuml.internals.interop import UnsupportedOnGPU
 from cuml.internals.validation import check_array
 
-__all__ = ("RandomForestRegressor", "RandomForestClassifier")
+__all__ = (
+    "IsolationForest",
+    "RandomForestClassifier",
+    "RandomForestRegressor",
+)
 
 
 class _RandomForestMixin:
@@ -88,3 +94,57 @@ def __iter__(self):
 
     def __getitem__(self, index):
         return self._call_method("__getitem__", index)
+
+
+class IsolationForest(ProxyBase):
+    _gpu_class = cuml.ensemble.IsolationForest
+    _other_attributes = frozenset(("_max_samples",))
+
+    @staticmethod
+    def _validate_input(X):
+        # Sparse inputs are handled by ProxyBase before dispatch. Fall back only
+        # for non-finite dense inputs, which scikit-learn supports but cuML does
+        # not, and preserve unrelated validation errors.
+        try:
+            check_array(
+                X, mem_type=None, order=None, ensure_2d=False, input_name="X"
+            )
+        except ValueError as exc:
+            message = str(exc)
+            if "contains NaN" in message or "contains infinity" in message:
+                raise UnsupportedOnGPU(message) from None
+            raise
+
+    def _gpu_fit(self, X, y=None, sample_weight=None):
+        self._validate_input(X)
+        if sample_weight is not None:
+            raise UnsupportedOnGPU("sample_weight is not supported")
+        return self._gpu.fit(X, y=y)
+
+    def _gpu_fit_predict(self, X, y=None, **kwargs):
+        # IsolationForest.fit_predict() doesn't declare sample_weight itself;
+        # it inherits OutlierMixin.fit_predict(self, X, y=None, **kwargs),
+        # which forwards kwargs straight to fit(). Match that signature here
+        # (rather than declaring sample_weight explicitly) so the proxy stays
+        # signature-compatible with the CPU method.
+        self._validate_input(X)
+        sample_weight = kwargs.pop("sample_weight", None)
+        if sample_weight is not None:
+            raise UnsupportedOnGPU("sample_weight is not supported")
+        if kwargs:
+            raise UnsupportedOnGPU(
+                "Additional fit parameters are not supported"
+            )
+        return self._gpu.fit_predict(X, y=y)
+
+    def _gpu_predict(self, X):
+        self._validate_input(X)
+        return self._gpu.predict(X)
+
+    def _gpu_decision_function(self, X):
+        self._validate_input(X)
+        return self._gpu.decision_function(X).astype(np.float64)
+
+    def _gpu_score_samples(self, X):
+        self._validate_input(X)
+        return self._gpu.score_samples(X).astype(np.float64)