Support Roots 2 (#219)

* Support Roots 2 * Add ForwardDiff rule for `find_alpha`
TuringLang · May 16, 2022 · b064e53 · b064e53 · devmotion · May 16, 2022
1 parent f987c1a
commit b064e53
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 2 deletions.
diff --git a/Project.toml b/Project.toml
@@ -1,6 +1,6 @@
 name = "Bijectors"
 uuid = "76274a88-744f-5084-9051-94815aaf08c4"
-version = "0.10.0"
+version = "0.10.1"
 
 [deps]
 ArgCheck = "dce04be8-c92d-5529-be00-80e4d2c0e197"
@@ -34,5 +34,5 @@ LogExpFunctions = "0.3.3"
 MappedArrays = "0.2.2, 0.3, 0.4"
 Reexport = "0.2, 1"
 Requires = "0.5, 1"
-Roots = "1.3.4"
+Roots = "1.3.4, 2"
 julia = "1.3"
diff --git a/src/compat/forwarddiff.jl b/src/compat/forwarddiff.jl
@@ -16,3 +16,26 @@ function jacobian(
 )
     return ForwardDiff.jacobian(b, x)
 end
+
+# Define forward-mode rule for ForwardDiff and don't trust support for ForwardDiff in Roots
+# https://github.com/JuliaMath/Roots.jl/issues/314
+function find_alpha(
+    wt_y::ForwardDiff.Dual{T,<:Real},
+    wt_u_hat::ForwardDiff.Dual{T,<:Real},
+    b::ForwardDiff.Dual{T,<:Real},
+) where {T}
+    # Compute primal
+    value_wt_y = ForwardDiff.value(wt_y)
+    value_wt_u_hat = ForwardDiff.value(wt_u_hat)
+    value_b = ForwardDiff.value(b)
+    Ω = find_alpha(value_wt_y, value_wt_u_hat, value_b)
+
+    # Compute derivative
+    partials_wt_y = ForwardDiff.partials(wt_y)
+    partials_wt_u_hat = ForwardDiff.partials(wt_u_hat)
+    partials_b = ForwardDiff.partials(b)
+    x = inv(1 + value_wt_u_hat * sech(Ω + value_b)^2)
+    ∂Ω = x * (partials_wt_y - tanh(Ω + value_b) * partials_wt_u_hat) + (x - 1) * partials_b
+
+    return ForwardDiff.Dual{T}(Ω, ∂Ω)
+end