fix: incorporate @return_on_nonfinite_val in LoopVectorization extension

Author: MilesCranmer

ext/DynamicExpressionsLoopVectorizationExt.jl

@@ -3,7 +3,7 @@ module DynamicExpressionsLoopVectorizationExt
 using LoopVectorization: @turbo
 using DynamicExpressions: AbstractExpressionNode
 using DynamicExpressions.UtilsModule: ResultOk, fill_similar
-using DynamicExpressions.EvaluateModule: @return_on_check, EvalOptions
+using DynamicExpressions.EvaluateModule: @return_on_nonfinite_val, EvalOptions
 import DynamicExpressions.EvaluateModule:
     deg1_eval,
     deg2_eval,
@@ -45,21 +45,21 @@ function deg1_l2_ll0_lr0_eval(
     cX::AbstractMatrix{T},
     op::F,
     op_l::F2,
-    ::EvalOptions{true},
+    eval_options::EvalOptions{true},
 ) where {T<:Number,F,F2}
     if tree.l.l.constant && tree.l.r.constant
         val_ll = tree.l.l.val
         val_lr = tree.l.r.val
-        @return_on_check val_ll cX
-        @return_on_check val_lr cX
+        @return_on_nonfinite_val(eval_options, val_ll, cX)
+        @return_on_nonfinite_val(eval_options, val_lr, cX)
         x_l = op_l(val_ll, val_lr)::T
-        @return_on_check x_l cX
+        @return_on_nonfinite_val(eval_options, x_l, cX)
         x = op(x_l)::T
-        @return_on_check x cX
+        @return_on_nonfinite_val(eval_options, x, cX)
         return ResultOk(fill_similar(x, cX, axes(cX, 2)), true)
     elseif tree.l.l.constant
         val_ll = tree.l.l.val
-        @return_on_check val_ll cX
+        @return_on_nonfinite_val(eval_options, val_ll, cX)
         feature_lr = tree.l.r.feature
         cumulator = similar(cX, axes(cX, 2))
         @turbo for j in axes(cX, 2)
@@ -71,7 +71,7 @@ function deg1_l2_ll0_lr0_eval(
     elseif tree.l.r.constant
         feature_ll = tree.l.l.feature
         val_lr = tree.l.r.val
-        @return_on_check val_lr cX
+        @return_on_nonfinite_val(eval_options, val_lr, cX)
         cumulator = similar(cX, axes(cX, 2))
         @turbo for j in axes(cX, 2)
             x_l = op_l(cX[feature_ll, j], val_lr)
@@ -97,15 +97,15 @@ function deg1_l1_ll0_eval(
     cX::AbstractMatrix{T},
     op::F,
     op_l::F2,
-    ::EvalOptions{true},
+    eval_options::EvalOptions{true},
 ) where {T<:Number,F,F2}
     if tree.l.l.constant
         val_ll = tree.l.l.val
-        @return_on_check val_ll cX
+        @return_on_nonfinite_val(eval_options, val_ll, cX)
         x_l = op_l(val_ll)::T
-        @return_on_check x_l cX
+        @return_on_nonfinite_val(eval_options, x_l, cX)
         x = op(x_l)::T
-        @return_on_check x cX
+        @return_on_nonfinite_val(eval_options, x, cX)
         return ResultOk(fill_similar(x, cX, axes(cX, 2)), true)
     else
         feature_ll = tree.l.l.feature
@@ -120,20 +120,23 @@ function deg1_l1_ll0_eval(
 end
 
 function deg2_l0_r0_eval(
-    tree::AbstractExpressionNode{T}, cX::AbstractMatrix{T}, op::F, ::EvalOptions{true}
+    tree::AbstractExpressionNode{T},
+    cX::AbstractMatrix{T},
+    op::F,
+    eval_options::EvalOptions{true},
 ) where {T<:Number,F}
     if tree.l.constant && tree.r.constant
         val_l = tree.l.val
-        @return_on_check val_l cX
+        @return_on_nonfinite_val(eval_options, val_l, cX)
         val_r = tree.r.val
-        @return_on_check val_r cX
+        @return_on_nonfinite_val(eval_options, val_r, cX)
         x = op(val_l, val_r)::T
-        @return_on_check x cX
+        @return_on_nonfinite_val(eval_options, x, cX)
         return ResultOk(fill_similar(x, cX, axes(cX, 2)), true)
     elseif tree.l.constant
         cumulator = similar(cX, axes(cX, 2))
         val_l = tree.l.val
-        @return_on_check val_l cX
+        @return_on_nonfinite_val(eval_options, val_l, cX)
         feature_r = tree.r.feature
         @turbo for j in axes(cX, 2)
             x = op(val_l, cX[feature_r, j])
@@ -144,7 +147,7 @@ function deg2_l0_r0_eval(
         cumulator = similar(cX, axes(cX, 2))
         feature_l = tree.l.feature
         val_r = tree.r.val
-        @return_on_check val_r cX
+        @return_on_nonfinite_val(eval_options, val_r, cX)
         @turbo for j in axes(cX, 2)
             x = op(cX[feature_l, j], val_r)
             cumulator[j] = x
@@ -168,11 +171,11 @@ function deg2_l0_eval(
     cumulator::AbstractVector{T},
     cX::AbstractArray{T},
     op::F,
-    ::EvalOptions{true},
+    eval_options::EvalOptions{true},
 ) where {T<:Number,F}
     if tree.l.constant
         val = tree.l.val
-        @return_on_check val cX
+        @return_on_nonfinite_val(eval_options, val, cX)
         @turbo for j in eachindex(cumulator)
             x = op(val, cumulator[j])
             cumulator[j] = x
@@ -193,11 +196,11 @@ function deg2_r0_eval(
     cumulator::AbstractVector{T},
     cX::AbstractArray{T},
     op::F,
-    ::EvalOptions{true},
+    eval_options::EvalOptions{true},
 ) where {T<:Number,F}
     if tree.r.constant
         val = tree.r.val
-        @return_on_check val cX
+        @return_on_nonfinite_val(eval_options, val, cX)
         @turbo for j in eachindex(cumulator)
             x = op(cumulator[j], val)
             cumulator[j] = x