From 95744b1867c551903434679eaa5334cd4d97216a Mon Sep 17 00:00:00 2001 From: MilesCranmerBot Date: Sat, 14 Feb 2026 13:45:48 +0000 Subject: [PATCH 1/7] fix: preserve torch export dims for column-vector inputs --- pysr/export_torch.py | 10 +++++++++- pysr/test/test_torch.py | 33 +++++++++++++++++++++++++++++++++ 2 files changed, 42 insertions(+), 1 deletion(-) diff --git a/pysr/export_torch.py b/pysr/export_torch.py index ef9aeaeb8..47ad2efb6 100644 --- a/pysr/export_torch.py +++ b/pysr/export_torch.py @@ -179,10 +179,18 @@ def __repr__(self): return f"{type(self).__name__}(expression={self._expression_string})" def forward(self, X): + if X.dim() == 1: + X = X.unsqueeze(-1) if self._selection is not None: X = X[:, self._selection] + if X.dim() == 1: + X = X.unsqueeze(-1) + preserve_2d_output = X.dim() == 2 and X.shape[1] == 1 symbols = {symbol: X[:, i] for i, symbol in enumerate(self.symbols_in)} - return self._node(symbols) + output = self._node(symbols) + if preserve_2d_output and output.dim() == 1: + output = output.unsqueeze(-1) + return output SingleSymPyModule = _SingleSymPyModule diff --git a/pysr/test/test_torch.py b/pysr/test/test_torch.py index e66e7d960..24e4efa62 100644 --- a/pysr/test/test_torch.py +++ b/pysr/test/test_torch.py @@ -192,6 +192,39 @@ def test_issue_656(self): decimal=3, ) + def test_issue_571_single_feature_shape(self): + x = sympy.symbols("x") + m = pysr.export_torch.sympy2torch(x + 1, [x]) + X = self.torch.randn(32, 1) + y = m(X) + self.assertEqual(tuple(y.shape), (32, 1)) + np.testing.assert_almost_equal( + y.detach().numpy().flatten(), + (X[:, 0] + 1).detach().numpy(), + decimal=6, + ) + + def test_issue_571_composition(self): + x = sympy.symbols("x") + a, b = sympy.symbols("a b") + m1 = pysr.export_torch.sympy2torch(x + 1, [x]) + m2 = pysr.export_torch.sympy2torch(2 * x, [x]) + m3 = pysr.export_torch.sympy2torch(a + b, [a, b]) + + X = self.torch.randn(32, 1) + y1 = m1(X[:, 0]) + y2 = m2(X[:, 0]) + self.assertEqual(tuple(y1.shape), (32, 1)) + self.assertEqual(tuple(y2.shape), (32, 1)) + + stacked = self.torch.cat([y1, y2], dim=1) + y3 = m3(stacked) + np.testing.assert_almost_equal( + y3.detach().numpy(), + (3 * X[:, 0] + 1).detach().numpy(), + decimal=6, + ) + def test_constant_arguments(self): # Test that functions with constant arguments work correctly # Regression test for https://github.com/MilesCranmer/PySR/issues/656 From 685daa190b17f6829a05ac4fd8e30562038d856c Mon Sep 17 00:00:00 2001 From: MilesCranmerBot Date: Sat, 14 Feb 2026 13:50:52 +0000 Subject: [PATCH 2/7] fix(export_torch): require 2D inputs and keep column-vector outputs --- pysr/export_torch.py | 27 +++++++++++++++++++-------- pysr/test/test_torch.py | 11 +++++++---- 2 files changed, 26 insertions(+), 12 deletions(-) diff --git a/pysr/export_torch.py b/pysr/export_torch.py index 47ad2efb6..d23a3c4b2 100644 --- a/pysr/export_torch.py +++ b/pysr/export_torch.py @@ -179,17 +179,28 @@ def __repr__(self): return f"{type(self).__name__}(expression={self._expression_string})" def forward(self, X): - if X.dim() == 1: - X = X.unsqueeze(-1) + if X.dim() != 2: + raise ValueError( + "Expected a 2D input tensor `X` with shape (L, nfeatures)." + ) + if self._selection is not None: X = X[:, self._selection] - if X.dim() == 1: - X = X.unsqueeze(-1) - preserve_2d_output = X.dim() == 2 and X.shape[1] == 1 - symbols = {symbol: X[:, i] for i, symbol in enumerate(self.symbols_in)} + + if X.dim() != 2: + raise ValueError( + "Expected a 2D input tensor `X` with shape (L, nfeatures)." + ) + + preserve_2d_output = X.shape[1] == 1 + symbols = { + symbol: X[:, i : i + 1] for i, symbol in enumerate(self.symbols_in) + } output = self._node(symbols) - if preserve_2d_output and output.dim() == 1: - output = output.unsqueeze(-1) + + if not preserve_2d_output and output.dim() == 2 and output.shape[1] == 1: + output = output.squeeze(-1) + return output SingleSymPyModule = _SingleSymPyModule diff --git a/pysr/test/test_torch.py b/pysr/test/test_torch.py index 24e4efa62..1dac9cb41 100644 --- a/pysr/test/test_torch.py +++ b/pysr/test/test_torch.py @@ -176,7 +176,7 @@ def test_avoid_simplification(self): rng = np.random.RandomState(0) X = rng.randn(10, 1) np.testing.assert_almost_equal( - m(torch.tensor(X)).detach().numpy(), + m(torch.tensor(X)).detach().numpy().flatten(), np.square(np.exp(np.sign(0.44796443))) + 1.5 * X[:, 0], decimal=3, ) @@ -187,7 +187,7 @@ def test_issue_656(self): m = pysr.export_torch.sympy2torch(E_plus_x1, ["x1"]) X = np.random.randn(10, 1) np.testing.assert_almost_equal( - m(self.torch.tensor(X)).detach().numpy(), + m(self.torch.tensor(X)).detach().numpy().flatten(), np.exp(1) + X[:, 0], decimal=3, ) @@ -212,11 +212,14 @@ def test_issue_571_composition(self): m3 = pysr.export_torch.sympy2torch(a + b, [a, b]) X = self.torch.randn(32, 1) - y1 = m1(X[:, 0]) - y2 = m2(X[:, 0]) + y1 = m1(X) + y2 = m2(X) self.assertEqual(tuple(y1.shape), (32, 1)) self.assertEqual(tuple(y2.shape), (32, 1)) + with self.assertRaises(ValueError): + m1(X[:, 0]) + stacked = self.torch.cat([y1, y2], dim=1) y3 = m3(stacked) np.testing.assert_almost_equal( From 7bf64da4283b2d1aed7d011392c82a44f039e357 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sat, 14 Feb 2026 13:51:08 +0000 Subject: [PATCH 3/7] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- pysr/export_torch.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/pysr/export_torch.py b/pysr/export_torch.py index d23a3c4b2..ae6a406fe 100644 --- a/pysr/export_torch.py +++ b/pysr/export_torch.py @@ -198,7 +198,11 @@ def forward(self, X): } output = self._node(symbols) - if not preserve_2d_output and output.dim() == 2 and output.shape[1] == 1: + if ( + not preserve_2d_output + and output.dim() == 2 + and output.shape[1] == 1 + ): output = output.squeeze(-1) return output From 88b3d281ffc85dfb03baea8a3ce1580b5955d0e8 Mon Sep 17 00:00:00 2001 From: MilesCranmerBot Date: Sat, 14 Feb 2026 14:41:42 +0000 Subject: [PATCH 4/7] test: use public sympy2torch API and cover dims --- pysr/test/test_torch.py | 38 +++++++++++++++++++++++++++----------- 1 file changed, 27 insertions(+), 11 deletions(-) diff --git a/pysr/test/test_torch.py b/pysr/test/test_torch.py index 1dac9cb41..cd0714e97 100644 --- a/pysr/test/test_torch.py +++ b/pysr/test/test_torch.py @@ -172,7 +172,7 @@ def test_avoid_simplification(self): feature_names_in=["x1"], extra_sympy_mappings={"square": lambda x: x**2}, ) - m = pysr.export_torch.sympy2torch(ex, ["x1"]) + m = sympy2torch(ex, ["x1"]) rng = np.random.RandomState(0) X = rng.randn(10, 1) np.testing.assert_almost_equal( @@ -184,7 +184,7 @@ def test_avoid_simplification(self): def test_issue_656(self): # Should correctly map numeric symbols to floats E_plus_x1 = sympy.exp(1) + sympy.symbols("x1") - m = pysr.export_torch.sympy2torch(E_plus_x1, ["x1"]) + m = sympy2torch(E_plus_x1, ["x1"]) X = np.random.randn(10, 1) np.testing.assert_almost_equal( m(self.torch.tensor(X)).detach().numpy().flatten(), @@ -194,7 +194,7 @@ def test_issue_656(self): def test_issue_571_single_feature_shape(self): x = sympy.symbols("x") - m = pysr.export_torch.sympy2torch(x + 1, [x]) + m = sympy2torch(x + 1, [x]) X = self.torch.randn(32, 1) y = m(X) self.assertEqual(tuple(y.shape), (32, 1)) @@ -204,12 +204,24 @@ def test_issue_571_single_feature_shape(self): decimal=6, ) + def test_issue_571_multifeature_output_is_1d(self): + x, y = sympy.symbols("x y") + m = sympy2torch(x + y, [x, y]) + X = self.torch.randn(32, 2) + out = m(X) + self.assertEqual(tuple(out.shape), (32,)) + np.testing.assert_almost_equal( + out.detach().numpy(), + (X[:, 0] + X[:, 1]).detach().numpy(), + decimal=6, + ) + def test_issue_571_composition(self): x = sympy.symbols("x") a, b = sympy.symbols("a b") - m1 = pysr.export_torch.sympy2torch(x + 1, [x]) - m2 = pysr.export_torch.sympy2torch(2 * x, [x]) - m3 = pysr.export_torch.sympy2torch(a + b, [a, b]) + m1 = sympy2torch(x + 1, [x]) + m2 = sympy2torch(2 * x, [x]) + m3 = sympy2torch(a + b, [a, b]) X = self.torch.randn(32, 1) y1 = m1(X) @@ -217,9 +229,6 @@ def test_issue_571_composition(self): self.assertEqual(tuple(y1.shape), (32, 1)) self.assertEqual(tuple(y2.shape), (32, 1)) - with self.assertRaises(ValueError): - m1(X[:, 0]) - stacked = self.torch.cat([y1, y2], dim=1) y3 = m3(stacked) np.testing.assert_almost_equal( @@ -228,6 +237,13 @@ def test_issue_571_composition(self): decimal=6, ) + def test_issue_571_reject_1d_input(self): + x = sympy.symbols("x") + m = sympy2torch(x + 1, [x]) + X = self.torch.randn(32, 1) + with self.assertRaises(ValueError): + m(X[:, 0]) + def test_constant_arguments(self): # Test that functions with constant arguments work correctly # Regression test for https://github.com/MilesCranmer/PySR/issues/656 @@ -239,14 +255,14 @@ def test_constant_arguments(self): ] for expr, expected in test_cases: - m = pysr.export_torch.sympy2torch(expr, []) + m = sympy2torch(expr, []) result = m(self.torch.randn(10, 1)) np.testing.assert_almost_equal(result.item(), expected, decimal=3) # Test with variables: sqrt(2) * x x = sympy.symbols("x") expr = sympy.sqrt(2) * x - m = pysr.export_torch.sympy2torch(expr, [x]) + m = sympy2torch(expr, [x]) X = np.random.randn(10, 1) np.testing.assert_almost_equal( m(self.torch.tensor(X)).detach().numpy().flatten(), From c75962405215524804113b476e64fac9e0af3a4e Mon Sep 17 00:00:00 2001 From: MilesCranmerBot Date: Sat, 14 Feb 2026 15:10:49 +0000 Subject: [PATCH 5/7] test: cover selection handling for torch export --- pysr/test/test_torch.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/pysr/test/test_torch.py b/pysr/test/test_torch.py index cd0714e97..88fd683b3 100644 --- a/pysr/test/test_torch.py +++ b/pysr/test/test_torch.py @@ -244,6 +244,20 @@ def test_issue_571_reject_1d_input(self): with self.assertRaises(ValueError): m(X[:, 0]) + def test_issue_571_selection_list_keeps_2d(self): + x = sympy.symbols("x") + m = sympy2torch(x + 1, [x], selection=[0]) + X = self.torch.randn(32, 2) + out = m(X) + self.assertEqual(tuple(out.shape), (32, 1)) + + def test_issue_571_reject_int_selection(self): + x = sympy.symbols("x") + m = sympy2torch(x + 1, [x], selection=0) + X = self.torch.randn(32, 2) + with self.assertRaises(ValueError): + m(X) + def test_constant_arguments(self): # Test that functions with constant arguments work correctly # Regression test for https://github.com/MilesCranmer/PySR/issues/656 From 78dfc8c187a63cd456ab429a326ca210d6582020 Mon Sep 17 00:00:00 2001 From: MilesCranmerBot Date: Sat, 14 Feb 2026 15:53:36 +0000 Subject: [PATCH 6/7] test: add docstrings for #571 regression tests --- pysr/test/test_torch.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/pysr/test/test_torch.py b/pysr/test/test_torch.py index 88fd683b3..fd003902d 100644 --- a/pysr/test/test_torch.py +++ b/pysr/test/test_torch.py @@ -193,6 +193,7 @@ def test_issue_656(self): ) def test_issue_571_single_feature_shape(self): + """Issue #571: 1-feature torch module preserves (L, 1) output shape.""" x = sympy.symbols("x") m = sympy2torch(x + 1, [x]) X = self.torch.randn(32, 1) @@ -205,6 +206,7 @@ def test_issue_571_single_feature_shape(self): ) def test_issue_571_multifeature_output_is_1d(self): + """Issue #571: multi-feature torch modules keep 1D outputs (L,) by default.""" x, y = sympy.symbols("x y") m = sympy2torch(x + y, [x, y]) X = self.torch.randn(32, 2) @@ -217,6 +219,7 @@ def test_issue_571_multifeature_output_is_1d(self): ) def test_issue_571_composition(self): + """Issue #571: composing 1-feature modules into a 2-feature module works.""" x = sympy.symbols("x") a, b = sympy.symbols("a b") m1 = sympy2torch(x + 1, [x]) @@ -238,6 +241,7 @@ def test_issue_571_composition(self): ) def test_issue_571_reject_1d_input(self): + """Issue #571: torch module rejects 1D inputs (expects (L, nfeatures)).""" x = sympy.symbols("x") m = sympy2torch(x + 1, [x]) X = self.torch.randn(32, 1) @@ -245,6 +249,7 @@ def test_issue_571_reject_1d_input(self): m(X[:, 0]) def test_issue_571_selection_list_keeps_2d(self): + """Issue #571: selection=[i] keeps (L, 1) shape after feature selection.""" x = sympy.symbols("x") m = sympy2torch(x + 1, [x], selection=[0]) X = self.torch.randn(32, 2) @@ -252,6 +257,7 @@ def test_issue_571_selection_list_keeps_2d(self): self.assertEqual(tuple(out.shape), (32, 1)) def test_issue_571_reject_int_selection(self): + """Issue #571: selection that collapses to 1D should raise (selection=0).""" x = sympy.symbols("x") m = sympy2torch(x + 1, [x], selection=0) X = self.torch.randn(32, 2) From a342033fef38897bdbd3b1bd815ce2157569dc9e Mon Sep 17 00:00:00 2001 From: MilesCranmerBot Date: Sat, 14 Feb 2026 16:40:32 +0000 Subject: [PATCH 7/7] fix(export_torch): drop redundant dim check after selection --- pysr/export_torch.py | 5 ----- pysr/test/test_torch.py | 8 -------- 2 files changed, 13 deletions(-) diff --git a/pysr/export_torch.py b/pysr/export_torch.py index ae6a406fe..fc27053df 100644 --- a/pysr/export_torch.py +++ b/pysr/export_torch.py @@ -187,11 +187,6 @@ def forward(self, X): if self._selection is not None: X = X[:, self._selection] - if X.dim() != 2: - raise ValueError( - "Expected a 2D input tensor `X` with shape (L, nfeatures)." - ) - preserve_2d_output = X.shape[1] == 1 symbols = { symbol: X[:, i : i + 1] for i, symbol in enumerate(self.symbols_in) diff --git a/pysr/test/test_torch.py b/pysr/test/test_torch.py index fd003902d..80e5e5f7f 100644 --- a/pysr/test/test_torch.py +++ b/pysr/test/test_torch.py @@ -256,14 +256,6 @@ def test_issue_571_selection_list_keeps_2d(self): out = m(X) self.assertEqual(tuple(out.shape), (32, 1)) - def test_issue_571_reject_int_selection(self): - """Issue #571: selection that collapses to 1D should raise (selection=0).""" - x = sympy.symbols("x") - m = sympy2torch(x + 1, [x], selection=0) - X = self.torch.randn(32, 2) - with self.assertRaises(ValueError): - m(X) - def test_constant_arguments(self): # Test that functions with constant arguments work correctly # Regression test for https://github.com/MilesCranmer/PySR/issues/656