Using DynamicDiff within CustomLossFunction

[gm89uk]

I'm just trying whether I can use DynamicDiff operator within a custom loss function and comparing that to eval_diff_tree_array

function loss_fnc(tree, dataset::Dataset{T,L}, options, idx) where {T,L}    
    # Extract data for the given indices
    X = idx === nothing ? dataset.X : view(dataset.X,:, idx)
    y = idx === nothing ? dataset.y : view(dataset.y, idx)
    newExp = D(tree, 2)
    grad, complete = eval_tree_array(newExp, X, options)
    if !complete
        return L(Inf)
    end
   # Monotonicity check
    if !all(h -> h <= 0, grad)
        return 1f09
    end
    #avoids second eval if monotonicity breached. 
    prediction, complete = eval_tree_array(tree, X, options)
    if !complete
        return L(Inf)
    end    
    #...
)

However, tree is DynamicExpressions.NodeModule.Node and I'm not sure how to extract AbstractExpression from tree. This would avoid any automatic differentiation and only evaluates the data if the constraints on grad are not breached.

Thank you

Edit: To clarify, I'm comparing against:

function loss_fnc(tree, dataset::Dataset{T,L}, options, idx) where {T,L}    
    # Extract data for the given indices
    X = idx === nothing ? dataset.X : view(dataset.X,:, idx)
    y = idx === nothing ? dataset.y : view(dataset.y, idx)
    prediction, grad, complete = eval_diff_tree_array(tree, X, options, 2)
    if !complete
        return L(Inf)
    end
    if !all(h -> h <= 0, grad)
        return 1f09
    end
    #...
)

I think if I get the first one working, it should be much quicker.

[MilesCranmer]

Thanks. Sadly there’s no way to do this at the moment because loss_function operates on the extracted tree, rather than the expression object. In the short term maybe we need a loss_function_expression to specify it as operating on the expression or something.

[gm89uk]

Thank you Miles, would be great.
Using D would avoid Zygote all together, as you'd just need eval_tree_array twice rather than a eval_diff_tree_array.

A very quick experiment of simulating two eval_tree_array:

function loss_fnc(tree, dataset::Dataset{T,L}, options, idx) where {T,L}    
    # Extract data for the given indices
    X = idx === nothing ? dataset.X : view(dataset.X,:, idx)
    y = idx === nothing ? dataset.y : view(dataset.y, idx)
    prediction, complete = eval_tree_array(tree, X, options)
    if !complete
        return L(Inf)
    end. 
    prediction, complete = eval_tree_array(tree, X, options)
    if !complete
        return L(Inf)
    end    
    #...
)

vs a single eval_diff_tree_array

function loss_fnc(tree, dataset::Dataset{T,L}, options, idx) where {T,L}    
    # Extract data for the given indices
    X = idx === nothing ? dataset.X : view(dataset.X,:, idx)
    y = idx === nothing ? dataset.y : view(dataset.y, idx)
    prediction, grad, complete = eval_diff_tree_array(tree, X, options, 2)
    if !complete
        return L(Inf)
    end
    #...
)

The first code is 3 times faster with my data! So it could be a way to significantly speed up eval_diff_tree_array by utilising DynamicDiff in the future.