From 762ff772f32692604b995721530a74624247e9c7 Mon Sep 17 00:00:00 2001
From: Kristoffer Carlsson <kcarlsson89@gmail.com>
Date: Mon, 12 Oct 2026 21:32:47 +0200
Subject: [PATCH] Reduce the cost of NaN-safe mode by checking the scalar
 factor once

In NaN-safe mode every partial was masked with `iszero(partial) ? zero(y) : y`.
The mask only changes the result when the scalar factor is non-finite (for
multiplication) or zero/NaN (for division), so check the factor once per tuple
operation and only take the masked path in that case. `_isfinite_all` recurses
into nested `Dual`s so Hessians also take the fast path.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 src/dual.jl     |  2 ++
 src/partials.jl | 38 +++++++++++++++++++++++++++++---------
 2 files changed, 31 insertions(+), 9 deletions(-)

diff --git a/src/dual.jl b/src/dual.jl
index 6d13dec3..1edfd88e 100644
--- a/src/dual.jl
+++ b/src/dual.jl
@@ -416,6 +416,8 @@ end
 
 Base.iszero(x::Dual) = iszero(value(x)) && iszero(partials(x))  # shortcut, equivalent to x == zero(x)
 
+@inline _isfinite_all(x::Dual) = _isfinite_all(value(x)) & all(_isfinite_all, partials(x).values)
+
 for pred in [:isequal, :(==)]
     @eval begin
         @define_binary_dual_op(
diff --git a/src/partials.jl b/src/partials.jl
index 49f8364b..cb4eef12 100644
--- a/src/partials.jl
+++ b/src/partials.jl
@@ -85,15 +85,27 @@ Base.convert(::Type{Partials{N,V}}, partials::Partials{N,V}) where {N,V} = parti
 @inline Base.:*(x::Real, partials::Partials) = partials*x
 
 @inline function Base.:*(partials::Partials, x::Real)
-    return Partials(scale_tuple(partials.values, x))
+    if _mul_needs_mask(x)
+        return Partials(scale_tuple(_mul_partial, partials.values, x))
+    else
+        return Partials(scale_tuple(*, partials.values, x))
+    end
 end
 
 @inline function Base.:/(partials::Partials, x::Real)
-    return Partials(div_tuple_by_scalar(partials.values, x))
+    if _div_needs_mask(x)
+        return Partials(div_tuple_by_scalar(_div_partial, partials.values, x))
+    else
+        return Partials(div_tuple_by_scalar(/, partials.values, x))
+    end
 end
 
 @inline function _mul_partials(a::Partials{N}, b::Partials{N}, x_a, x_b) where N
-    return Partials(mul_tuples(a.values, b.values, x_a, x_b))
+    if _mul_needs_mask(x_a) | _mul_needs_mask(x_b)
+        return Partials(mul_tuples(_mul_partial, a.values, b.values, x_a, x_b))
+    else
+        return Partials(mul_tuples(*, a.values, b.values, x_a, x_b))
+    end
 end
 
 @inline function _div_partials(a::Partials, b::Partials, aval, bval)
@@ -117,11 +129,19 @@ if NANSAFE_MODE_ENABLED
         y = partial / x
         return iszero(partial) ? zero(y) : y
     end
+    @inline _mul_needs_mask(x::Real) = !_isfinite_all(x)
+    @inline _div_needs_mask(x::Real) = true
+    @inline _div_needs_mask(x::Base.IEEEFloat) = !isfinite(x) | iszero(x)
 else
     @inline _mul_partial(partial::Real, x::Real) = partial * x
     @inline _div_partial(partial::Real, x::Real) = partial / x
+    @inline _mul_needs_mask(x::Real) = false
+    @inline _div_needs_mask(x::Real) = false
 end
 
+# Whether `x` and, for dual numbers, all of its partials are finite
+@inline _isfinite_all(x::Real) = isfinite(x)
+
 # edge cases where N == 0 #
 #-------------------------#
 
@@ -199,12 +219,12 @@ end
     return tupexpr(i -> :(rand(V)), N)
 end
 
-@generated function scale_tuple(tup::NTuple{N}, x) where N
-    return tupexpr(i -> :(_mul_partial(tup[$i], x)), N)
+@generated function scale_tuple(f::F, tup::NTuple{N}, x) where {F,N}
+    return tupexpr(i -> :(f(tup[$i], x)), N)
 end
 
-@generated function div_tuple_by_scalar(tup::NTuple{N}, x) where N
-    return tupexpr(i -> :(_div_partial(tup[$i], x)), N)
+@generated function div_tuple_by_scalar(f::F, tup::NTuple{N}, x) where {F,N}
+    return tupexpr(i -> :(f(tup[$i], x)), N)
 end
 
 @generated function add_tuples(a::NTuple{N}, b::NTuple{N})  where N
@@ -219,8 +239,8 @@ end
     return tupexpr(i -> :(-tup[$i]), N)
 end
 
-@generated function mul_tuples(a::NTuple{N}, b::NTuple{N}, afactor, bfactor) where N
-    return tupexpr(i -> :(_mul_partial(a[$i], afactor) + _mul_partial(b[$i], bfactor)), N)
+@generated function mul_tuples(f::F, a::NTuple{N}, b::NTuple{N}, afactor, bfactor) where {F,N}
+    return tupexpr(i -> :(f(a[$i], afactor) + f(b[$i], bfactor)), N)
 end
 
 ###################