Add code to reproduce error

Author: c-w-feldmann

test_extras/test_chemprop/chemprop_test_utils/default_models.py

@@ -1,5 +1,7 @@
 """Functions for creating default chemprop models."""
 
+from typing import Any
+
 from molpipeline.estimators.chemprop import ChempropModel, ChempropNeuralFP
 from molpipeline.estimators.chemprop.component_wrapper import (
     MPNN,
@@ -28,16 +30,26 @@ def get_binary_classification_mpnn() -> MPNN:
     return mpnn
 
 
-def get_neural_fp_encoder() -> ChempropNeuralFP:
+def get_neural_fp_encoder(
+    init_kwargs: dict[str, Any] | None = None,
+) -> ChempropNeuralFP:
     """Get the Chemprop model.
 
+    Parameters
+    ----------
+    init_kwargs : dict[str, Any], optional
+        Additional keyword arguments to pass to `ChempropNeuralFP` during initialization.
+
     Returns
     -------
     ChempropNeuralFP
         The Chemprop model.
     """
     mpnn = get_binary_classification_mpnn()
-    chemprop_model = ChempropNeuralFP(model=mpnn, lightning_trainer__accelerator="cpu")
+    init_kwargs = init_kwargs or {}
+    chemprop_model = ChempropNeuralFP(
+        model=mpnn, lightning_trainer__accelerator="cpu", **init_kwargs
+    )
     return chemprop_model
 
 

test_extras/test_chemprop/test_neural_fingerprint.py

@@ -7,8 +7,6 @@
 
 from molpipeline.estimators.chemprop.neural_fingerprint import ChempropNeuralFP
 from molpipeline.utils.json_operations import recursive_from_json, recursive_to_json
-
-# pylint: disable=relative-beyond-top-level
 from test_extras.test_chemprop.chemprop_test_utils.compare_models import compare_params
 from test_extras.test_chemprop.chemprop_test_utils.default_models import (
     get_neural_fp_encoder,
@@ -38,3 +36,10 @@ def test_output_type(self) -> None:
         """Test the output type."""
         chemprop_fp_encoder = get_neural_fp_encoder()
         self.assertEqual(chemprop_fp_encoder.output_type, "float")
+
+    def test_init_with_kwargs(self) -> None:
+        """Test the __init__ method with kwargs."""
+        init_kwargs = {"model__message_passing__depth": 4}
+        chemprop_fp_encoder = get_neural_fp_encoder(init_kwargs=init_kwargs)
+        deep_params = chemprop_fp_encoder.get_params(deep=True)
+        self.assertEqual(deep_params["model__message_passing__depth"], 4)