fix: mutation error for Zygote

Author: MilesCranmer

src/EvaluationHelpers.jl

@@ -1,11 +1,19 @@
 module EvaluationHelpersModule
 
+using ChainRulesCore: @non_differentiable
+
 import Base: adjoint
 import ..OperatorEnumModule: AbstractOperatorEnum, OperatorEnum, GenericOperatorEnum
 import ..NodeModule: AbstractExpressionNode
 import ..EvaluateModule: eval_tree_array
 import ..EvaluateDerivativeModule: eval_grad_tree_array
 
+function _set_nan!(out)
+    out .= convert(eltype(out), NaN)
+    return nothing
+end
+@non_differentiable _set_nan!(out)
+
 # Evaluation:
 """
     (tree::AbstractExpressionNode)(X, operators::OperatorEnum; kws...)
@@ -27,7 +35,7 @@ and triplets of operations for lower memory usage.
 """
 function (tree::AbstractExpressionNode)(X, operators::OperatorEnum; kws...)
     out, did_finish = eval_tree_array(tree, X, operators; kws...)
-    !did_finish && (out .= convert(eltype(out), NaN))
+    !did_finish && _set_nan!(out)
     return out
 end
 """
@@ -56,7 +64,7 @@ function _grad_evaluator(
     tree::AbstractExpressionNode, X, operators::OperatorEnum; variable=Val(true), kws...
 )
     _, grad, did_complete = eval_grad_tree_array(tree, X, operators; variable, kws...)
-    !did_complete && (grad .= convert(eltype(grad), NaN))
+    !did_complete && _set_nan!(grad)
     return grad
 end
 function _grad_evaluator(