BoundsError when using TemplateExpressionSpec with custom loss function

[yejinkim-repo]

Hi, I’m hitting a BoundsError whenever I use TemplateExpressionSpec with a simple setup (combine="g(s) + h(v)"). It looks like the template combine function is trying to access expressions[2] while only one sub-expression is present.

My Environment

PySR: 1.5.9
SymbolicRegression: 1.12.0

Code

import numpy as np
from pysr import jl, PySRRegressor, TemplateExpressionSpec

rng = np.random.default_rng(0)
x1 = rng.normal(size=200)
x2 = rng.normal(size=200)
X = np.column_stack([x1, x2])
y = np.sin(x1) + np.cos(x2)

# To verify custom loss operation
custom_objective = """
function my_custom_objective(tree, dataset::Dataset{T,L}, options) where {T,L}
    pred, flag = eval_tree_array(tree, dataset.X, options)
    !flag && return L(Inf)
    return sum((pred .- dataset.y).^2) / dataset.n
end
"""

template = TemplateExpressionSpec(
    expressions=["f", "g"],
    variable_names=["x1", "x2"],
    combine="f(x1) + g(x2)",
)

model = PySRRegressor(
    expression_spec=template,
    binary_operators=["+", "*"],
    unary_operators=["exp"],
    loss_function=custom_objective,
    variable_names=["x1", "x2"],
    niterations=100,
)

model.fit(X, y)

Error message

JuliaError: BoundsError: attempt to access 1-element Vector{ComposableExpression{Float32, Node{Float32}, @NamedTuple{operators::OperatorEnum{Tuple{typeof(+)}, Tuple{typeof(exp)}}, variable_names::Nothing, eval_options::EvalOptions{false, false, true, Nothing}}}} at index [2]
Stacktrace:
  [1] throw_boundserror(A::Vector{ComposableExpression{Float32, Node{Float32}, @NamedTuple{operators::OperatorEnum{Tuple{typeof(+)}, Tuple{typeof(exp)}}, variable_names::Nothing, eval_options::EvalOptions{false, false, true, Nothing}}}}, I::Tuple{Int64})
    @ Base ./essentials.jl:14
  [2] getindex
    @ ./essentials.jl:916 [inlined]
...

I don't know why this problem occurs and how to fix it, but if you have any good suggestions, please let me know.
Thanks a lot for your help!

[MilesCranmer]

Ah, you should use loss_function_expression, rather than loss_function. I think a warning is supposed to get printed suggesting this - does that show up?

[yejinkim-repo]

Thanks, using loss_function_expression did resolve the original issue — the code runs fine now with eval_tree_array.

However, when I try to extend the custom loss to also compute gradients, switching to eval_grad_tree_array(...; variable=true), I hit the same BoundsError again.

Do you think this means the tree isn’t being passed/expanded correctly into eval_grad_tree_array when using a TemplateExpressionSpec? Or is there a different recommended way to get input-variable gradients inside a custom loss?

[MilesCranmer]

For computing gradients in a template expression you could use the D operator? https://ai.damtp.cam.ac.uk/pysr/examples/#13-using-differential-operators

eval_grad_tree_array is not set up to work with template expressions. You could either use D within the template expression body (recommended) or, for advanced usage, you could use ForwardDiff (or Zygote/Mooncake/Enzyme) to compute a derivative of eval_tree_array directly.

In any case, it's probably a good idea to throw an error for both cases (much appreciated if you could post a dedicated issue for this, to help us track things)

[yejinkim-repo]

Thanks for the clarification!

Just to confirm my use case - I don’t want the gradients to appear in the template expression output. I just want to use computed gradients inside a custom loss function, while keeping the model output as f(x1) + g(x2).

Given that eval_grad_tree_array is not supported for templates, I will likely need to differentiate eval_tree_array directly with ForwardDiff (or Zygote/Mooncake/Enzyme) in order to compute partial derivatives with respect to x1 and x2.

I’ve started experimenting with ForwardDiff, but I’m still struggling to figure out how to differentiate the tree directly inside a custom function. I’ll open a new discussion to track the process of resolving this.

[MilesCranmer]

You can use template expressions as a convenient way of writing loss functions; there are some other discussions on this worth reading. For example, you could pass y and dy as “features” and then return abs(y - f(x)) + abs(dy - D(f, 1)(x)) from the template expressions; with elementwise_loss=(pred, target) -> pred