refactor: clean up use of is_valid throughout library

Author: MilesCranmer

ext/DynamicExpressionsBumperExt.jl

@@ -1,8 +1,9 @@
 module DynamicExpressionsBumperExt
 
 using Bumper: @no_escape, @alloc
-using DynamicExpressions: OperatorEnum, AbstractExpressionNode, tree_mapreduce
-using DynamicExpressions.UtilsModule: ResultOk, counttuple, is_bad_array
+using DynamicExpressions:
+    OperatorEnum, AbstractExpressionNode, tree_mapreduce, is_valid_array
+using DynamicExpressions.UtilsModule: ResultOk, counttuple
 
 import DynamicExpressions.ExtensionInterfaceModule:
     bumper_eval_tree_array, bumper_kern1!, bumper_kern2!
@@ -52,7 +53,7 @@ function dispatch_kerns!(operators, branch_node, cumulator, ::Val{turbo}) where
     cumulator.ok || return cumulator
 
     out = dispatch_kern1!(operators.unaops, branch_node.op, cumulator.x, Val(turbo))
-    return ResultOk(out, !is_bad_array(out))
+    return ResultOk(out, is_valid_array(out))
 end
 function dispatch_kerns!(
     operators, branch_node, cumulator1, cumulator2, ::Val{turbo}
@@ -63,7 +64,7 @@ function dispatch_kerns!(
     out = dispatch_kern2!(
         operators.binops, branch_node.op, cumulator1.x, cumulator2.x, Val(turbo)
     )
-    return ResultOk(out, !is_bad_array(out))
+    return ResultOk(out, is_valid_array(out))
 end
 
 @generated function dispatch_kern1!(unaops, op_idx, cumulator, ::Val{turbo}) where {turbo}

ext/DynamicExpressionsSymbolicUtilsExt.jl

@@ -4,7 +4,8 @@ using SymbolicUtils
 import DynamicExpressions.NodeModule:
     AbstractExpressionNode, Node, constructorof, DEFAULT_NODE_TYPE
 import DynamicExpressions.OperatorEnumModule: AbstractOperatorEnum
-import DynamicExpressions.UtilsModule: isgood, isbad, deprecate_varmap
+import DynamicExpressions.ValueInterfaceModule: is_valid
+import DynamicExpressions.UtilsModule: deprecate_varmap
 import DynamicExpressions.ExtensionInterfaceModule: node_to_symbolic, symbolic_to_node
 import DynamicExpressions: AbstractExpression, get_tree, get_operators
 
@@ -19,14 +20,14 @@ macro return_on_false(flag, retval)
     )
 end
 
-function isgood(x::SymbolicUtils.Symbolic)
+function is_valid(x::SymbolicUtils.Symbolic)
     return if SymbolicUtils.istree(x)
-        all(isgood.([SymbolicUtils.operation(x); SymbolicUtils.arguments(x)]))
+        all(is_valid.([SymbolicUtils.operation(x); SymbolicUtils.arguments(x)]))
     else
         true
     end
 end
-subs_bad(x) = isgood(x) ? x : Inf
+subs_bad(x) = is_valid(x) ? x : Inf
 
 function parse_tree_to_eqs(
     tree::AbstractExpressionNode{T},
@@ -197,7 +198,7 @@ function node_to_symbolic(
     variable_names = deprecate_varmap(variable_names, varMap, :node_to_symbolic)
     expr = subs_bad(parse_tree_to_eqs(tree, operators, index_functions))
     # Check for NaN and Inf
-    @assert isgood(expr) "The recovered equation contains NaN or Inf."
+    @assert is_valid(expr) "The recovered equation contains NaN or Inf."
     # Return if no variable_names is given
     variable_names === nothing && return expr
     # Create a substitution tuple
@@ -248,12 +249,12 @@ function multiply_powers(
     if nargs == 1
         l, complete = multiply_powers(args[1])
         @return_on_false complete eqn
-        @return_on_false isgood(l) eqn
+        @return_on_false is_valid(l) eqn
         return op(l), true
     elseif op == ^
         l, complete = multiply_powers(args[1])
         @return_on_false complete eqn
-        @return_on_false isgood(l) eqn
+        @return_on_false is_valid(l) eqn
         n = args[2]
         if typeof(n) <: Integer
             if n == 1
@@ -275,23 +276,23 @@ function multiply_powers(
     elseif nargs == 2
         l, complete = multiply_powers(args[1])
         @return_on_false complete eqn
-        @return_on_false isgood(l) eqn
+        @return_on_false is_valid(l) eqn
         r, complete2 = multiply_powers(args[2])
         @return_on_false complete2 eqn
-        @return_on_false isgood(r) eqn
+        @return_on_false is_valid(r) eqn
         return op(l, r), true
     else
         # return tree_mapreduce(multiply_powers, op, args)
         # ## reduce(op, map(multiply_powers, args))
         out = map(multiply_powers, args) #vector of tuples
         for i in 1:size(out, 1)
             @return_on_false out[i][2] eqn
-            @return_on_false isgood(out[i][1]) eqn
+            @return_on_false is_valid(out[i][1]) eqn
         end
         cumulator = out[1][1]
         for i in 2:size(out, 1)
             cumulator = op(cumulator, out[i][1])
-            @return_on_false isgood(cumulator) eqn
+            @return_on_false is_valid(cumulator) eqn
         end
         return cumulator, true
     end

src/Evaluate.jl

@@ -5,7 +5,7 @@ using DispatchDoctor: @unstable
 import ..NodeModule: AbstractExpressionNode, constructorof
 import ..StringsModule: string_tree
 import ..OperatorEnumModule: OperatorEnum, GenericOperatorEnum
-import ..UtilsModule: is_bad_array, fill_similar, counttuple, ResultOk
+import ..UtilsModule: fill_similar, counttuple, ResultOk
 import ..NodeUtilsModule: is_constant
 import ..ExtensionInterfaceModule: bumper_eval_tree_array, _is_loopvectorization_loaded
 import ..ValueInterfaceModule: is_valid, is_valid_array

src/EvaluateDerivative.jl

@@ -2,8 +2,9 @@ module EvaluateDerivativeModule
 
 import ..NodeModule: AbstractExpressionNode, constructorof
 import ..OperatorEnumModule: OperatorEnum
-import ..UtilsModule: is_bad_array, fill_similar, ResultOk2
-import ..NodeUtilsModule: count_scalar_constants, index_constant_nodes, NodeIndex
+import ..UtilsModule: fill_similar, ResultOk2
+import ..ValueInterfaceModule: is_valid_array
+import ..NodeUtilsModule: count_constant_nodes, index_constant_nodes, NodeIndex
 import ..EvaluateModule: deg0_eval, get_nuna, get_nbin, OPERATOR_LIMIT_BEFORE_SLOWDOWN
 import ..ExtensionInterfaceModule: _zygote_gradient
 
@@ -105,7 +106,7 @@ end
         end
         !result.ok && return result
         return ResultOk2(
-            result.x, result.dx, !(is_bad_array(result.x) || is_bad_array(result.dx))
+            result.x, result.dx, is_valid_array(result.x) && is_valid_array(result.dx)
         )
     end
 end
@@ -213,9 +214,9 @@ function eval_grad_tree_array(
     n_gradients = if variable_mode
         size(cX, 1)::Int
     elseif constant_mode
-        count_scalar_constants(tree)::Int
+        count_constant_nodes(tree)::Int
     elseif both_mode
-        size(cX, 1) + count_scalar_constants(tree)
+        size(cX, 1) + count_constant_nodes(tree)
     end
 
     result = if variable_mode
@@ -247,7 +248,7 @@ function eval_grad_tree_array(
     result = _eval_grad_tree_array(tree, n_gradients, index_tree, cX, operators, Val(mode))
     !result.ok && return result
     return ResultOk2(
-        result.x, result.dx, !(is_bad_array(result.x) || is_bad_array(result.dx))
+        result.x, result.dx, is_valid_array(result.x) && is_valid_array(result.dx)
     )
 end
 

src/Simplify.jl

@@ -3,7 +3,7 @@ module SimplifyModule
 import ..NodeModule: AbstractExpressionNode, constructorof, Node, copy_node, set_node!
 import ..NodeUtilsModule: tree_mapreduce, is_node_constant
 import ..OperatorEnumModule: AbstractOperatorEnum
-import ..UtilsModule: isbad, isgood
+import ..ValueInterfaceModule: is_valid
 
 _una_op_kernel(f::F, l::T) where {F,T} = f(l)
 _bin_op_kernel(f::F, l::T, r::T) where {F,T} = f(l, r)
@@ -109,13 +109,13 @@ end
 function combine_children!(operators, p::N, c::N...) where {T,N<:AbstractExpressionNode{T}}
     all(is_node_constant, c) || return p
     vals = map(n -> n.val, c)
-    all(isgood, vals) || return p
+    all(is_valid, vals) || return p
     out = if length(c) == 1
         _una_op_kernel(operators.unaops[p.op], vals...)
     else
         _bin_op_kernel(operators.binops[p.op], vals...)
     end
-    isgood(out) || return p
+    is_valid(out) || return p
     new_node = constructorof(N)(T; val=convert(T, out))
     set_node!(p, new_node)
     return p

src/Utils.jl

@@ -12,13 +12,6 @@ macro return_on_false2(flag, retval, retval2)
     )
 end
 
-# Fastest way to check for NaN in an array.
-# (due to optimizations in sum())
-is_bad_array(array) = !(isempty(array) || isfinite(sum(array)))
-isgood(x::T) where {T<:Number} = !(isnan(x) || !isfinite(x))
-isgood(x) = true
-isbad(x) = !isgood(x)
-
 """
     @memoize_on tree [postprocess] function my_function_on_tree(tree::AbstractExpressionNode)
         ...