Better test coverage, and bsub(::Dual) specialization.

Author: chriselrod

src/RecursiveTupleMath.jl

@@ -27,12 +27,12 @@ using ForwardDiff: ≺, @define_binary_dual_op, Dual
 @inline lt_fast(a, b) = a < b
 @inline lt_fast(a::Float64, b::Float64) = Base.lt_float_fast(a, b)
 @inline lt_fast(a::Float32, b::Float32) = Base.lt_float_fast(a, b)
-@inline le_fast(a, b) = a <= b
-@inline le_fast(a::Float64, b::Float64) = Base.le_float_fast(a, b)
-@inline le_fast(a::Float32, b::Float32) = Base.le_float_fast(a, b)
-
 @inline gt_fast(a, b) = lt_fast(b, a)
-@inline ge_fast(a, b) = le_fast(b, a)
+
+# @inline le_fast(a, b) = a <= b
+# @inline le_fast(a::Float64, b::Float64) = Base.le_float_fast(a, b)
+# @inline le_fast(a::Float32, b::Float32) = Base.le_float_fast(a, b)
+# @inline ge_fast(a, b) = le_fast(b, a)
 
 @inline badd(a, b) = Base.FastMath.add_fast(a, b)
 @inline bsub(a, b) = Base.FastMath.sub_fast(a, b)
@@ -47,8 +47,6 @@ for bf ∈ [:bmax, :bmin, :badd, :bsub, :bmul, :bdiv]
     # terminating case
     @inline $bf(::Number, ::Tuple{}) = ()
     @inline $bf(::Tuple{}, ::Number) = ()
-    @inline $bf(::Number, ::Nothing) = nothing
-    @inline $bf(::Nothing, ::Number) = nothing
 
     # broadcast
     @inline $bf(x::Number, y::StaticArray{S}) where {S} = SArray{S}($bf(x, Tuple(y)))
@@ -87,60 +85,60 @@ end
   ntuple(Returns(v), Val(D))
 end
 
-# @inline 
+@inline bsub(x::Dual{T}) where {T} = Dual{T}(bsub(x.value), bsub(x.partials.values))
 
 @define_binary_dual_op(
   RecursiveTupleMath.badd,
   begin
-    Base.@_inline_meta
+    @_inline
     Dual{Txy}(badd(x.value, y.value), badd(x.partials.values, y.partials.values))
   end,
   begin
-    Base.@_inline_meta
+    @_inline
     Dual{Tx}(badd(x.value, y), x.partials)
   end,
   begin
-    Base.@_inline_meta
+    @_inline
     Dual{Ty}(badd(x, y.value), y.partials.values)
   end,
 )
 @define_binary_dual_op(
   RecursiveTupleMath.bsub,
   begin
-    Base.@_inline_meta
+    @_inline
     Dual{Txy}(bsub(x.value, y.value), bsub(x.partials.values, y.partials.values))
   end,
   begin
-    Base.@_inline_meta
+    @_inline
     Dual{Tx}(bsub(x.value, y), x.partials)
   end,
   begin
-    Base.@_inline_meta
+    @_inline
     Dual{Ty}(bsub(x, y.value), bsub(y.partials.values))
   end
 )
 @define_binary_dual_op(
   RecursiveTupleMath.bmul,
   begin
-    Base.@_inline_meta
+    @_inline
     Dual{Txy}(
       bmul(x.value, y.value),
       badd(bmul(x.value, y.partials.values), bmul(x.partials.values, y.value)),
     )
   end,
   begin
-    Base.@_inline_meta
+    @_inline
     Dual{Tx}(bmul(x.value, y), bmul(x.partials.values, y))
   end,
   begin
-    Base.@_inline_meta
+    @_inline
     Dual{Ty}(bmul(x, y.value), bmul(x, y.partials.values))
   end
 )
 @define_binary_dual_op(
   RecursiveTupleMath.bdiv,
   begin
-    Base.@_inline_meta
+    @_inline
     Dual{Txy}(
       bdiv(x.value, y.value),
       bdiv(
@@ -150,11 +148,11 @@ end
     )
   end,
   begin
-    Base.@_inline_meta
+    @_inline
     Dual{Tx}(bdiv(x.value, y), bdiv(bmul(x.partials.values, y), bmul(y, y)))
   end,
   begin
-    Base.@_inline_meta
+    @_inline
     Dual{Ty}(
       bdiv(x, y.value),
       bdiv(bsub(bmul(x, y.partials.values)), bmul(y.value, y.value)),
@@ -164,15 +162,15 @@ end
 @define_binary_dual_op(
   RecursiveTupleMath.bmax,
   begin
-    Base.@_inline_meta
+    @_inline
     cmp = gt_fast(x.value, y.value)
     v = ifelse(cmp, x.value, y.value)
     bcmp = btuple(cmp, Val(length(x.partials)))
     p = map(ifelse, bcmp, x.partials.values, y.partials.values)
     Dual{Txy}(v, p)
   end,
   begin
-    Base.@_inline_meta
+    @_inline
     cmp = gt_fast(x.value, y)
     v = ifelse(cmp, x.value, y)
     bcmp = btuple(cmp, Val(length(x.partials)))
@@ -181,7 +179,7 @@ end
     Dual{Tx}(v, p)
   end,
   begin
-    Base.@_inline_meta
+    @_inline
     cmp = gt_fast(x, y.value)
     v = ifelse(cmp, x, y.value)
     bcmp = btuple(cmp, Val(length(y.partials)))
@@ -193,15 +191,15 @@ end
 @define_binary_dual_op(
   RecursiveTupleMath.bmin,
   begin
-    Base.@_inline_meta
+    @_inline
     cmp = lt_fast(x.value, y.value)
     v = ifelse(cmp, x.value, y.value)
     bcmp = btuple(cmp, Val(length(x.partials)))
     p = map(ifelse, bcmp, x.partials.values, y.partials.values)
     Dual{Txy}(v, p)
   end,
   begin
-    Base.@_inline_meta
+    @_inline
     cmp = lt_fast(x.value, y)
     v = ifelse(cmp, x.value, y)
     bcmp = btuple(cmp, Val(length(x.partials)))
@@ -210,7 +208,7 @@ end
     Dual{Tx}(v, p)
   end,
   begin
-    Base.@_inline_meta
+    @_inline
     cmp = lt_fast(x, y.value)
     v = ifelse(cmp, x, y.value)
     bcmp = btuple(cmp, Val(length(y.partials)))

test/runtests.jl

@@ -2,32 +2,59 @@ using Base: Forward
 using RecursiveTupleMath, StaticArrays, ForwardDiff
 using Test
 
-foo(x, y, z) = sum(min.(z .+ x, (x .- y)) ./ max.(z .- y, (x .* y)))
-bar(x, y, z) = sum(bdiv(bmin(badd(z, x), bsub(x, y)), bmax(bsub(z, y), bmul(x, y))))
+abc(x::Number) = x
+abc(x::SVector{3}) = (a = x[1], b = x[2], c = x[3])
 
-@testset "RecursiveTupleMath.jl" begin
-  x = rand()
-  xv = @SVector rand(3)
-  for y in (rand(), @SVector(rand(3))), z in (rand(), @SVector(rand(3)))
+foo0(x, y, z) = sum(min.(z .+ x, (x .- y)) ./ max.(z .- y, (x .* y)))
+bar0(x, y, z) = sum(bdiv(bmin(badd(z, x), bsub(x, y)), bmax(bsub(z, y), bmul(x, y))))
+
+foo1(x, y, z) = sum((min.(x, y) ./ max.(x, y)) .* z .- z)
+bar1(x, y, z) = sum(bsub(bmul(bdiv(bmin(x, y), bmax(x, y)), z), z))
 
-    @test ForwardDiff.derivative(x -> foo(x, y, z), x) ≈
-          ForwardDiff.derivative(x -> bar(x, y, z), x)
-    @test ForwardDiff.derivative(x -> foo(y, x, z), x) ≈
-          ForwardDiff.derivative(x -> bar(y, x, z), x)
-    @test ForwardDiff.derivative(x -> foo(y, z, x), x) ≈
-          ForwardDiff.derivative(x -> bar(y, z, x), x)
+foo2(x, y, z) = sum((.-(x) ./ y) .+ (z ./ x))
+bar2(x, y, z) = sum(badd(bdiv(bsub(x), y), bdiv(z, x)))
 
-    @test ForwardDiff.gradient(x -> foo(x, y, z), xv) ≈
-          ForwardDiff.gradient(x -> bar(x, y, z), xv)
-    @test ForwardDiff.gradient(x -> foo(y, x, z), xv) ≈
-          ForwardDiff.gradient(x -> bar(y, x, z), xv)
-    @test ForwardDiff.gradient(x -> foo(y, z, x), xv) ≈
-          ForwardDiff.gradient(x -> bar(y, z, x), xv)
-    @test ForwardDiff.hessian(x -> foo(x, y, z), xv) ≈
-          ForwardDiff.hessian(x -> bar(x, y, z), xv)
-    @test ForwardDiff.hessian(x -> foo(y, x, z), xv) ≈
-          ForwardDiff.hessian(x -> bar(y, x, z), xv)
-    @test ForwardDiff.hessian(x -> foo(y, z, x), xv) ≈
-          ForwardDiff.hessian(x -> bar(y, z, x), xv)
+struct Cycle{F,Y,Z}
+  f::F
+  y::Y
+  z::Z
+  n::Int
+end
+Cycle(f::F, y::Y, z::Z) where {F,Y,Z} = Cycle(f, y, z, 0)
+Cycle(c::Cycle, n::Int) = Cycle(c.f, c.y, c.z, n)
+function (f::Cycle)(x)
+  y = f.y
+  z = f.z
+  if f.n == 0
+    f.f(x, y, z)
+  elseif f.n == 1
+    f.f(y, x, z)
+  else
+    f.f(y, z, x)
+  end
+end
+
+@testset "RecursiveTupleMath.jl" begin
+  for T ∈ (Float64, Float32)
+    x = rand(T)
+    xv = @SVector rand(T, 3)
+    for (f, b) ∈ ((foo0, bar0), (foo1, bar1), (foo2, bar2))
+      for y in (rand(T), @SVector(rand(T, 3))), z in (rand(T), @SVector(rand(T, 3)))
+        let g = Cycle(f, y, z), h = Cycle(b, y, z), habc = Cycle(b, abc(y), abc(z))
+          for n = 1:3
+            @test ForwardDiff.derivative(g, x) ≈ ForwardDiff.derivative(h, x)
+            @test ForwardDiff.gradient(g, xv) ≈
+                  ForwardDiff.gradient(h, xv) ≈
+                  ForwardDiff.gradient(habc ∘ abc, xv)
+            @test ForwardDiff.hessian(g, xv) ≈
+                  ForwardDiff.hessian(h, xv) ≈
+                  ForwardDiff.hessian(habc ∘ abc, xv)
+            g = Cycle(g, n)
+            h = Cycle(h, n)
+            habc = Cycle(habc, n)
+          end
+        end
+      end
+    end
   end
 end