Parametrize NumVal with type

BayesWitnesses · Apr 4, 2020 · 242165e · 242165e
1 parent e13bea3
commit 242165e
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 2 deletions.
diff --git a/m2cgen/assemblers/boosting.py b/m2cgen/assemblers/boosting.py
@@ -134,7 +134,7 @@ def _assemble_tree(self, tree):
         if "leaf" in tree:
             return ast.NumVal(tree["leaf"])
 
-        threshold = ast.NumVal(np.float32(tree["split_condition"]))
+        threshold = ast.NumVal(tree["split_condition"], dtype=np.float32)
         split = tree["split"]
         feature_idx = self._feature_name_to_idx.get(split, split)
         feature_ref = ast.FeatureRef(feature_idx)

diff --git a/m2cgen/ast.py b/m2cgen/ast.py
@@ -1,3 +1,4 @@
+import numpy as np
 from enum import Enum
 
 
@@ -30,7 +31,9 @@ class NumExpr(Expr):
 
 
 class NumVal(NumExpr):
-    def __init__(self, value):
+    def __init__(self, value, dtype=None):
+        if dtype:
+            value = dtype(value)
         self.value = value
 
     def __str__(self):