Merge pull request #131 from SymbolicML/fix-zygote-mutation

fix: move non_differentiable to special module for JET masking

Author: MilesCranmer

src/EvaluationHelpers.jl

@@ -8,11 +8,11 @@ import ..NodeModule: AbstractExpressionNode
 import ..EvaluateModule: eval_tree_array
 import ..EvaluateDerivativeModule: eval_grad_tree_array
 
+# Needs to be special function so we can declare it non-differentiable to Zygote
 function _set_nan!(out)
     out .= convert(eltype(out), NaN)
     return nothing
 end
-@non_differentiable _set_nan!(out)
 
 # Evaluation:
 """

src/NonDifferentiableDeclarations.jl

@@ -6,13 +6,15 @@ import ..NodeModule: AbstractExpressionNode, AbstractNode
 import ..NodeUtilsModule: tree_mapreduce
 import ..ExpressionModule:
     AbstractExpression, get_operators, get_variable_names, _validate_input
+import ..EvaluationHelpersModule: _set_nan!
 
 #! format: off
 @non_differentiable tree_mapreduce(f::Function, op::Function, tree::AbstractNode, result_type::Type)
 @non_differentiable tree_mapreduce(f::Function, f_branch::Function, op::Function, tree::AbstractNode, result_type::Type)
 @non_differentiable get_operators(ex::Union{AbstractExpression,AbstractExpressionNode}, operators::Union{AbstractOperatorEnum,Nothing})
 @non_differentiable get_variable_names(ex::AbstractExpression, variable_names::Union{AbstractVector{<:AbstractString},Nothing})
 @non_differentiable _validate_input(ex::AbstractExpression, X, operators::Union{AbstractOperatorEnum,Nothing})
+@non_differentiable _set_nan!(::Any)
 #! format: on
 
 end