Fix edgecase when expression is a constant

MilesCranmer · Apr 26, 2022 · 10bac39 · 10bac39
1 parent 7a792a8
commit 10bac39
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/pysr/sr.py b/pysr/sr.py
@@ -271,12 +271,13 @@ def __repr__(self):
         return f"PySRFunction(X=>{self._sympy})"
 
     def __call__(self, X):
+        expected_shape = (X.shape[0],)
         if isinstance(X, pd.DataFrame):
             # Lambda function takes as argument:
-            return self._lambda(**{k: X[k].values for k in X.columns})
+            return self._lambda(**{k: X[k].values for k in X.columns}) * np.ones(expected_shape)
         elif self._selection is not None:
-            return self._lambda(*X[:, self._selection].T)
-        return self._lambda(*X.T)
+            return self._lambda(*X[:, self._selection].T) * np.ones(expected_shape)
+        return self._lambda(*X.T) * np.ones(expected_shape)
 
 
 def _get_julia_project(julia_project):