Sparse hyperplanes have different dimensionality
@@ -637,6 +637,13 @@ def recursive_flatten(tree, hyperplanes, offsets,
def flatten_tree(tree, leaf_size):
n_nodes = num_nodes(tree)
n_leaves = num_leaves(tree)
+ if len(tree.hyperplane.shape) > 1:
+ # sparse case
+ hyperplanes = np.zeros((n_nodes,
+ tree.hyperplane.shape[0],
+ tree.hyperplane.shape[1]),
+ dtype=np.float64)
+ else:
hyperplanes = np.zeros((n_nodes, tree.hyperplane.shape[0]),
dtype=np.float64)
offsets = np.zeros(n_nodes, dtype=np.float64)
