Only check spox installed as depends on onnx

Quantco · Jul 24, 2024 · de5d495 · de5d495
1 parent a025e78
commit de5d495
Show file tree

Hide file tree

Showing 5 changed files with 0 additions and 9 deletions.
diff --git a/metalearners/_utils.py b/metalearners/_utils.py
@@ -548,7 +548,6 @@ def infer_probabilities_output(model) -> tuple[int, str]:
 def infer_input_dict(model) -> dict:
     """Returns a dict where the keys are the input names of the model and the values are
     ``spox.Var`` with the corresponding shape and type."""
-    check_onnx_installed()
     check_spox_installed()
     from spox import Tensor, Var, argument
 

diff --git a/metalearners/drlearner.py b/metalearners/drlearner.py
@@ -20,7 +20,6 @@
     _ScikitModel,
 )
 from metalearners._utils import (
-    check_onnx_installed,
     check_spox_installed,
     clip_element_absolute_value_to_epsilon,
     copydoc,
@@ -417,7 +416,6 @@ def build_onnx(self, models: Mapping[str, Sequence], output_name: str = "tau"):
         * ``"treatment_model"``
         """
         warning_experimental_feature("build_onnx")
-        check_onnx_installed()
         check_spox_installed()
         import spox.opset.ai.onnx.v21 as op
         from onnx.checker import check_model

diff --git a/metalearners/rlearner.py b/metalearners/rlearner.py
@@ -11,7 +11,6 @@
 
 from metalearners._typing import Matrix, OosMethod, Scoring, Vector
 from metalearners._utils import (
-    check_onnx_installed,
     check_spox_installed,
     clip_element_absolute_value_to_epsilon,
     copydoc,
@@ -537,7 +536,6 @@ def build_onnx(self, models: Mapping[str, Sequence], output_name: str = "tau"):
         * ``"treatment_model"``
         """
         warning_experimental_feature("build_onnx")
-        check_onnx_installed()
         check_spox_installed()
         import spox.opset.ai.onnx.v21 as op
         from onnx.checker import check_model

diff --git a/metalearners/tlearner.py b/metalearners/tlearner.py
@@ -9,7 +9,6 @@
 
 from metalearners._typing import Matrix, OosMethod, Scoring, Vector
 from metalearners._utils import (
-    check_onnx_installed,
     check_spox_installed,
     copydoc,
     index_matrix,
@@ -162,7 +161,6 @@ def build_onnx(self, models: Mapping[str, Sequence], output_name: str = "tau"):
         * ``"variant_outcome_model"``
         """
         warning_experimental_feature("build_onnx")
-        check_onnx_installed()
         check_spox_installed()
         import spox.opset.ai.onnx.v21 as op
         from onnx.checker import check_model

diff --git a/metalearners/xlearner.py b/metalearners/xlearner.py
@@ -10,7 +10,6 @@
 
 from metalearners._typing import Matrix, OosMethod, Scoring, Vector
 from metalearners._utils import (
-    check_onnx_installed,
     check_spox_installed,
     copydoc,
     get_one,
@@ -452,7 +451,6 @@ def build_onnx(self, models: Mapping[str, Sequence], output_name: str = "tau"):
         * ``"treatment_effect_model"``
         """
         warning_experimental_feature("build_onnx")
-        check_onnx_installed()
         check_spox_installed()
         import spox.opset.ai.onnx.v21 as op
         from onnx.checker import check_model