Allow y to be vector-valued?

[patrick-kidger]

It would be nice to allow y to be vector valued, essentially as a (parallelised?) shortcut for something like the following:

expressions = []
for index in range(y.shape[-1]):
    yi = y[:, index]
    expression_i = pysr.pysr(x, yi)
    expressions.append(expression_i)

This links to (and is actually motivated by) my suggestion in #32: as_pytorch could then return a single SymPyModule wrapping a list of all expressions, rather than returning several SymPyModules for each one, which would then have to be composed together in the above for loop.

[MilesCranmer]

Definitely agree. Actually it could be even nicer to just have the SymbolicRegression.jl backend loop over each y-component and return a list of equations at each complexity. Will think more about this.

[MilesCranmer]

Just added capabilities for multi-output (and multiple equations returned). SymbolicRegression.jl will now return a list of equations for each output.

This should be simple enough to set up for just importing into PySR. The trickier part is the JAX and PyTorch backend! I'm not sure how to approach - whether to continue outputting a different Torch/JAX model for each equation, or try to merge? It's probably overcomplicated to merge...

[patrick-kidger]

Is the PyTorch backend (you mean exporting via sympytorch?) actually available? I don't see it based on a quick glance through the code.

For sympytorch at least it's possible to pass multiple equations to the SymPyModule, i.e. vector values are automatically supported.

[MilesCranmer]

The PyTorch backend hasn't been implemented yet; it's next on my todo list!

That is nice that vector equations are supported! It doesn't assume the equations are all the same, but for different parameter choices, right? They can be completely different equations?

[patrick-kidger]

Yep, completely different equations.

[MilesCranmer]

Actually on second thought I'm not sure it will work due to the API of PySR: it will output a list of equations for each feature of y. When JAX/Torch export is turned on, it will generate a separate function/module for each expression. But I think the user will want to filter equations for each feature of y separately before making a choice of equation, and then fusing them.
I don't see an easy way to do this otherwise? I also don't know if there'd be a speedup since the expressions are assumed to be different, and no operations are re-used?

[patrick-kidger]

Hmm. I see your point. Perhaps only offer an export as SymPy, and just document that there are external libraries to convert SymPy into PyTorch/JAX. Then the end user can do whatever they like to select their sympy expression from the various possibilities.