Indirectly learning multiple subexpressions where main expression is a surrogate model

[Jes-s-L]

Hello! I am currently trying to solve a problem of the form $f(T) = g(a(T), b(T), ...)$ to learn the expressions of parameters $a(T), b(T), ...$ but where $g$ is a model represented by a neural network with inputs $a(T), b(T), ...$ etc.

Currently I have been able to successfully solve for a single parameter, i.e. $f(T) = g(a(T))$, by calling the NN representing $g$ from an onnx file in a custom loss function to evaluate $f(T)$ with each $a(T)$ expression, and calculate MSE against some imported known data for $f(T)$ from a csv.

However, I am unsure if it is possible to extend this to multiple parameters. My understanding is that when an array is passed to model.fit to fit multiple functions against a single output, the loss function and expression trees for each output are evaluated separately and therefore the losses cannot be combined.

Could I get some advice on whether this is a reasonable approach and if there is a method to combine the evaluation of multiple outputs in one loss function?

Thank you in advance!

[MilesCranmer]

Hello, have you looked at TemplateExpressionSpec? There are some examples in the docs (https://ai.damtp.cam.ac.uk/pysr/v1.5.9/examples.html#_11-expression-specifications) and various more complex examples discussed in other threads

[Jes-s-L]

Ah sorry I did not realise I could pass an entire julia snippet. I'll give that a go. Thank you!