Remove ADBackend and ADBijector (#242)

* remove AD backend related stuff

* removed usages of ADBijector

* removed inclusion of deleted file adbijector

* removed further references to ADBijector

* fix for ordered

* removed Identity which should been removed before

* make bijector for TransformedDistribution slightly nicer

* bumped minor version

Project.toml

@@ -1,7 +1,6 @@
 name = "Bijectors"
 uuid = "76274a88-744f-5084-9051-94815aaf08c4"
-version = "0.11.1"
-
+version = "0.12.0"
 
 [deps]
 ArgCheck = "dce04be8-c92d-5529-be00-80e4d2c0e197"

README.md

@@ -33,7 +33,7 @@ The following table lists mathematical operations for a bijector and the corresp
 
 In this table, `b` denotes a `Bijector`, `J(b, x)` denotes the Jacobian of `b` evaluated at `x`, `b_*` denotes the [push-forward](https://www.wikiwand.com/en/Pushforward_measure) of `p` by `b`, and `x ∼ p` denotes `x` sampled from the distribution with density `p`.
 
-The "Automatic" column in the table refers to whether or not you are required to implement the feature for a custom `Bijector`. "AD" refers to the fact that it can be implemented "automatically" using automatic differentiation, i.e. `ADBijector`.
+The "Automatic" column in the table refers to whether or not you are required to implement the feature for a custom `Bijector`. "AD" refers to the fact that this can be implemented "automatically" using automatic differentiation, e.g. ForwardDiff.jl.
 
 ## Functions
 

src/Bijectors.jl

@@ -64,11 +64,9 @@ export  TransformDistribution,
         logabsdetjac!,
         logabsdetjacinv,
         Bijector,
-        ADBijector,
         Inverse,
         Stacked,
         stack,
-        Identity,
         bijector,
         transformed,
         UnivariateTransformed,

src/bijectors/ordered.jl

@@ -17,10 +17,10 @@ Return a `Distribution` whose support are ordered vectors, i.e., vectors with in
 This transformation is currently only supported for otherwise unconstrained distributions.
 """
 function ordered(d::ContinuousMultivariateDistribution)
-    if !isa(bijector(d), Identity)
+    if bijector(d) !== identity
         throw(ArgumentError("ordered transform is currently only supported for unconstrained distributions."))
     end
-    return Bijectors.transformed(d, OrderedBijector())
+    return transformed(d, OrderedBijector())
 end
 
 with_logabsdet_jacobian(b::OrderedBijector, x) = transform(b, x), logabsdetjac(b, x)

src/bijectors/stacked.jl

@@ -16,7 +16,7 @@ where `bs[i]::Bijector` is applied to `x[ranges[i]]::UnitRange{Int}`.
 # Examples
 ```
 b1 = Logit(0.0, 1.0)
-b2 = Identity()
+b2 = identity
 b = stack(b1, b2)
 b([0.0, 1.0]) == [b1(0.0), 1.0]  # => true
 ```

src/compat/distributionsad.jl

@@ -9,22 +9,22 @@ using Distributions: AbstractMvLogNormal
 bijector(::TuringDirichlet) = SimplexBijector()
 bijector(::TuringWishart) = PDBijector()
 bijector(::TuringInverseWishart) = PDBijector()
-bijector(::TuringScalMvNormal) = Identity()
-bijector(::TuringDiagMvNormal) = Identity()
-bijector(::TuringDenseMvNormal) = Identity()
+bijector(::TuringScalMvNormal) = identity
+bijector(::TuringDiagMvNormal) = identity
+bijector(::TuringDenseMvNormal) = identity
 
 bijector(d::FillVectorOfUnivariate{Continuous}) = bijector(d.v.value)
 bijector(d::FillMatrixOfUnivariate{Continuous}) = up1(bijector(d.dists.value))
-bijector(d::MatrixOfUnivariate{Discrete}) = Identity()
+bijector(d::MatrixOfUnivariate{Discrete}) = identity
 bijector(d::MatrixOfUnivariate{Continuous}) = TruncatedBijector(_minmax(d.dists)...)
-bijector(d::VectorOfMultivariate{Discrete}) = Identity()
+bijector(d::VectorOfMultivariate{Discrete}) = identity
 for T in (:VectorOfMultivariate, :FillVectorOfMultivariate)
     @eval begin
-        bijector(d::$T{Continuous, <:MvNormal}) = Identity()
-        bijector(d::$T{Continuous, <:TuringScalMvNormal}) = Identity()
-        bijector(d::$T{Continuous, <:TuringDiagMvNormal}) = Identity()
-        bijector(d::$T{Continuous, <:TuringDenseMvNormal}) = Identity()
-        bijector(d::$T{Continuous, <:MvNormalCanon}) = Identity()
+        bijector(d::$T{Continuous, <:MvNormal}) = identity
+        bijector(d::$T{Continuous, <:TuringScalMvNormal}) = identity
+        bijector(d::$T{Continuous, <:TuringDiagMvNormal}) = identity
+        bijector(d::$T{Continuous, <:TuringDenseMvNormal}) = identity
+        bijector(d::$T{Continuous, <:MvNormalCanon}) = identity
         bijector(d::$T{Continuous, <:AbstractMvLogNormal}) = Log()
         bijector(d::$T{Continuous, <:SimplexDistribution}) = SimplexBijector()
         bijector(d::$T{Continuous, <:TuringDirichlet}) = SimplexBijector()

src/compat/forwarddiff.jl

@@ -3,20 +3,6 @@ import .ForwardDiff
 _eps(::Type{<:ForwardDiff.Dual{<:Any, Real}}) = _eps(Real)
 _eps(::Type{<:ForwardDiff.Dual{<:Any, <:Integer}}) = _eps(Real)
 
-# AD implementations
-function jacobian(
-    b::Union{<:ADBijector{<:ForwardDiffAD}, Inverse{<:ADBijector{<:ForwardDiffAD}}},
-    x::Real
-)
-    return ForwardDiff.derivative(b, x)
-end
-function jacobian(
-    b::Union{<:ADBijector{<:ForwardDiffAD}, Inverse{<:ADBijector{<:ForwardDiffAD}}},
-    x::AbstractVector{<:Real}
-)
-    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(

src/compat/reversediff.jl

@@ -5,8 +5,7 @@ using ..ReverseDiff: ReverseDiff, @grad, value, track, TrackedReal, TrackedVecto
 using Requires, LinearAlgebra
 
 using ..Bijectors: Elementwise, SimplexBijector, maphcat, simplex_link_jacobian, 
-    simplex_invlink_jacobian, simplex_logabsdetjac_gradient, ADBijector, 
-    ReverseDiffAD, Inverse
+    simplex_invlink_jacobian, simplex_logabsdetjac_gradient, Inverse
 import ..Bijectors: _eps, logabsdetjac, _logabsdetjac_scale, _simplex_bijector, 
     _simplex_inv_bijector, replace_diag, jacobian, getpd, lower, 
     _inv_link_chol_lkj, _link_chol_lkj, _transform_ordered, _transform_inverse_ordered,
@@ -17,20 +16,6 @@ import ChainRulesCore
 using Compat: eachcol
 using Distributions: LocationScale
 
-# AD implementations
-function jacobian(
-    b::Union{<:ADBijector{<:ReverseDiffAD}, Inverse{<:ADBijector{<:ReverseDiffAD}}},
-    x::Real
-)
-    return ReverseDiff.gradient(x -> b(x[1]), [x])[1]
-end
-function jacobian(
-    b::Union{<:ADBijector{<:ReverseDiffAD}, Inverse{<:ADBijector{<:ReverseDiffAD}}},
-    x::AbstractVector{<:Real}
-)
-    return ReverseDiff.jacobian(b, x)
-end
-
 _eps(::Type{<:TrackedReal{T}}) where {T} = _eps(T)
 function Base.minimum(d::LocationScale{<:TrackedReal})
     m = minimum(d.ρ)

src/compat/tracker.jl

@@ -12,8 +12,7 @@ using ..Tracker: Tracker,
                  param
 
 import ..Bijectors
-using ..Bijectors: Elementwise, SimplexBijector, ADBijector,
-    TrackerAD, Inverse, Stacked
+using ..Bijectors: Elementwise, SimplexBijector, Inverse, Stacked
 
 import ChainRulesCore
 import LogExpFunctions
@@ -49,21 +48,6 @@ function Base.maximum(d::LocationScale{<:TrackedReal})
     end
 end
 
-# AD implementations
-function Bijectors.jacobian(
-    b::Union{<:ADBijector{<:TrackerAD}, Inverse{<:ADBijector{<:TrackerAD}}},
-    x::Real
-)
-    return data(Tracker.gradient(b, x)[1])
-end
-function Bijectors.jacobian(
-    b::Union{<:ADBijector{<:TrackerAD}, Inverse{<:ADBijector{<:TrackerAD}}},
-    x::AbstractVector{<:Real}
-)
-    # We extract `data` so that we don't return a `Tracked` type
-    return data(Tracker.jacobian(b, x))
-end
-
 # implementations for Shift bijector
 function Bijectors._logabsdetjac_shift(a::TrackedReal, x::Real, ::Val{0})
     return tracker_shift_logabsdetjac(a, x, Val(0))

src/compat/zygote.jl

@@ -32,18 +32,6 @@ end
 end
 
 # AD implementations
-function jacobian(
-    b::Union{<:ADBijector{<:ZygoteAD}, Inverse{<:ADBijector{<:ZygoteAD}}},
-    x::Real
-)
-    return Zygote.gradient(b, x)[1]
-end
-function jacobian(
-    b::Union{<:ADBijector{<:ZygoteAD}, Inverse{<:ADBijector{<:ZygoteAD}}},
-    x::AbstractVector{<:Real}
-)
-    return Zygote.jacobian(b, x)
-end
 @adjoint function _logabsdetjac_scale(a::Real, x::Real, ::Val{0})
     return _logabsdetjac_scale(a, x, Val(0)), Δ -> (inv(a) .* Δ, nothing, nothing)
 end

src/interface.jl

@@ -18,31 +18,6 @@ elementwise(f) = Base.Fix1(broadcast, f)
 # the way to go.
 elementwise(f::ComposedFunction) = ComposedFunction(elementwise(f.outer), elementwise(f.inner))
 
-#######################################
-# AD stuff "extracted" from Turing.jl #
-#######################################
-
-abstract type ADBackend end
-struct ForwardDiffAD <: ADBackend end
-struct ReverseDiffAD <: ADBackend end
-struct TrackerAD <: ADBackend end
-struct ZygoteAD <: ADBackend end
-
-const ADBACKEND = Ref(:forwarddiff)
-setadbackend(backend_sym::Symbol) = setadbackend(Val(backend_sym))
-setadbackend(::Val{:forwarddiff}) = ADBACKEND[] = :forwarddiff
-setadbackend(::Val{:reversediff}) = ADBACKEND[] = :reversediff
-setadbackend(::Val{:tracker}) = ADBACKEND[] = :tracker
-setadbackend(::Val{:zygote}) = ADBACKEND[] = :zygote
-
-ADBackend() = ADBackend(ADBACKEND[])
-ADBackend(T::Symbol) = ADBackend(Val(T))
-ADBackend(::Val{:forwarddiff}) = ForwardDiffAD
-ADBackend(::Val{:reversediff}) = ReverseDiffAD
-ADBackend(::Val{:tracker}) = TrackerAD
-ADBackend(::Val{:zygote}) = ZygoteAD
-ADBackend(::Val) = error("The requested AD backend is not available. Make sure to load all required packages.")
-
 ######################
 # Bijector interface #
 ######################
@@ -197,12 +172,8 @@ Just an alias for `logabsdetjac(inverse(b), y)`.
 logabsdetjacinv(b, y) = logabsdetjac(inverse(b), y)
 
 ##############################
-# Example bijector: Identity #
+# Example bijector: identity #
 ##############################
-Identity() = identity
-
-# Here we don't need to separate between batched version and non-batched, and so
-# we can just overload `transform`, etc. directly.
 transform(::typeof(identity), x) = copy(x)
 transform!(::typeof(identity), x, y) = copy!(y, x)
 
@@ -213,7 +184,6 @@ logabsdetjac!(::typeof(identity), x, logjac) = logjac
 # Bijectors includes #
 ######################
 # General
-include("bijectors/adbijector.jl")
 include("bijectors/composed.jl")
 include("bijectors/stacked.jl")
 

src/transformed_distribution.jl

@@ -35,11 +35,14 @@ transformed(d) = transformed(d, bijector(d))
 
 Returns the constrained-to-unconstrained bijector for distribution `d`.
 """
-bijector(td::TransformedDistribution) = bijector(td.dist) ∘ inverse(td.transform)
-bijector(d::DiscreteUnivariateDistribution) = Identity()
-bijector(d::DiscreteMultivariateDistribution) = Identity()
+function bijector(td::TransformedDistribution)
+    b = bijector(td.dist)
+    return b === identity ? inverse(td.transform) : b ∘ inverse(td.transform)
+end
+bijector(d::DiscreteUnivariateDistribution) = identity
+bijector(d::DiscreteMultivariateDistribution) = identity
 bijector(d::ContinuousUnivariateDistribution) = TruncatedBijector(minimum(d), maximum(d))
-bijector(d::Product{Discrete}) = Identity()
+bijector(d::Product{Discrete}) = identity
 function bijector(d::Product{Continuous})
     return TruncatedBijector(_minmax(d.v)...)
 end
@@ -52,8 +55,8 @@ end
     end
 end
 
-bijector(d::Normal) = Identity()
-bijector(d::Distributions.AbstractMvNormal) = Identity()
+bijector(d::Normal) = identity
+bijector(d::Distributions.AbstractMvNormal) = identity
 bijector(d::Distributions.AbstractMvLogNormal) = elementwise(log)
 bijector(d::PositiveDistribution) = elementwise(log)
 bijector(d::SimplexDistribution) = SimplexBijector()

test/bijectors/ordered.jl

@@ -23,7 +23,7 @@ end
     @test d_ordered.dist === d
     @test d_ordered.transform isa OrderedBijector
     y = randn(5)
-    x = inv(bijector(d_ordered))(y)
+    x = inverse(bijector(d_ordered))(y)
     @test issorted(x)
 
     d = Product(fill(Normal(), 5))