Merge pull request #52 from tpapp/tp/zygote-tests

Get rid of RealVector, views, test Zygote.

Author: tpapp

CHANGELOG.md

@@ -2,6 +2,8 @@
 
 - use new LogDensityProblems interface
 
+- rewrite internals to work better with AD (especially Zygote)
+
 # 0.3.4
 
 - make `inverse(::ArrayTransform)` accept `AbstractArray`

Manifest.toml

@@ -130,7 +130,7 @@ uuid = "d96e819e-fc66-5662-9728-84c9c7592b0a"
 version = "0.11.0"
 
 [[Pkg]]
-deps = ["Dates", "LibGit2", "Markdown", "Printf", "REPL", "Random", "SHA", "UUIDs"]
+deps = ["Dates", "LibGit2", "Libdl", "Logging", "Markdown", "Printf", "REPL", "Random", "SHA", "UUIDs"]
 uuid = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 
 [[Printf]]

Project.toml

@@ -10,6 +10,7 @@ ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
 Parameters = "d96e819e-fc66-5662-9728-84c9c7592b0a"
+Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 
 [compat]
@@ -24,6 +25,7 @@ Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
-test = ["Flux", "LogDensityProblems", "OffsetArrays", "Random", "ReverseDiff", "StaticArrays", "Test"]
+test = ["Flux", "LogDensityProblems", "OffsetArrays", "Random", "ReverseDiff", "StaticArrays", "Test", "Zygote"]

src/aggregation.jl

@@ -54,37 +54,36 @@ function as(::Type{Matrix}, args...)
     t
 end
 
-function transform_with(flag::LogJacFlag, t::ArrayTransform, x::RealVector)
+function transform_with(flag::LogJacFlag, t::ArrayTransform, x, index::T) where {T}
     @unpack transformation, dims = t
+    # NOTE not using index increments as that somehow breaks type inference
     d = dimension(transformation)
-    I = reshape(range(firstindex(x); length = prod(dims), step = d), dims)
-    yℓ = map(i -> transform_with(flag, transformation, view_into(x, i, d)), I)
+    𝐼 = reshape(range(index; length = prod(dims), step = d), dims)
+    yℓ = map(index -> ((y, ℓ, _) = transform_with(flag, transformation, x, index); (y, ℓ)), 𝐼)
     ℓz = logjac_zero(flag, extended_eltype(x))
-    first.(yℓ), isempty(yℓ) ? ℓz : ℓz + sum(last, yℓ)
+    first.(yℓ), isempty(yℓ) ? ℓz : ℓz + sum(last, yℓ), index
 end
 
-function transform_with(flag::LogJacFlag, t::ArrayTransform{Identity}, x::RealVector)
+function transform_with(flag::LogJacFlag, t::ArrayTransform{Identity}, x, index)
     # TODO use version below when https://github.com/FluxML/Flux.jl/issues/416 is fixed
     # y = reshape(copy(x), t.dims)
-    y = reshape(map(identity, x), t.dims)
-    y, logjac_zero(flag, extended_eltype(x))
+    index′ = index+dimension(t)
+    y = reshape(map(identity, x[index:(index′-1)]), t.dims)
+    y, logjac_zero(flag, extended_eltype(x)), index′
 end
 
-inverse_eltype(t::ArrayTransform, x::AbstractArray) =
+function inverse_eltype(t::ArrayTransform, x::AbstractArray)
     inverse_eltype(t.transformation, first(x)) # FIXME shortcut
+end
 
-function inverse!(x::RealVector,
-                  transformation_array::ArrayTransform,
-                  y::AbstractArray)
+function inverse_at!(x::AbstractVector, index, transformation_array::ArrayTransform,
+                     y::AbstractArray)
     @unpack transformation, dims = transformation_array
     @argcheck size(y) == dims
-    index = firstindex(x)
-    d = dimension(transformation)
     for elt in vec(y)
-        inverse!(view_into(x, index, d), transformation, elt)
-        index += d;
+        index = inverse_at!(x, index, transformation, elt)
     end
-    x
+    index
 end
 
 ####
@@ -152,24 +151,24 @@ $(SIGNATURES)
 Helper function for transforming tuples. Used internally, to help type inference. Use via
 `transfom_tuple`.
 """
-_transform_tuple(flag::LogJacFlag, x::RealVector, index, ::Tuple{}) =
-    (), logjac_zero(flag, extended_eltype(x))
+_transform_tuple(flag::LogJacFlag, x::AbstractVector, index, ::Tuple{}) =
+    (), logjac_zero(flag, extended_eltype(x)), index
 
-function _transform_tuple(flag::LogJacFlag, x::RealVector, index, ts)
+function _transform_tuple(flag::LogJacFlag, x::AbstractVector, index, ts)
     tfirst = first(ts)
-    d = dimension(tfirst)
-    yfirst, ℓfirst = transform_with(flag, tfirst, view_into(x, index, d))
-    yrest, ℓrest = _transform_tuple(flag, x, index + d, Base.tail(ts))
-    (yfirst, yrest...), ℓfirst + ℓrest
+    yfirst, ℓfirst, index′ = transform_with(flag, tfirst, x, index)
+    yrest, ℓrest, index′′ = _transform_tuple(flag, x, index′, Base.tail(ts))
+    (yfirst, yrest...), ℓfirst + ℓrest, index′′
 end
 
 """
 $(SIGNATURES)
 
 Helper function for tuple transformations.
 """
-transform_tuple(flag::LogJacFlag, tt::NTransforms, x::RealVector) =
-    _transform_tuple(flag, x, firstindex(x), tt)
+function transform_tuple(flag::LogJacFlag, tt::NTransforms, x, index)
+    _transform_tuple(flag, x, index, tt)
+end
 
 """
 $(SIGNATURES)
@@ -189,39 +188,37 @@ Helper function for inverting tuples of transformations. Used internally.
 
 *Performs no argument validation, caller should do this.*
 """
-function _inverse!_tuple(x::RealVector, ts::NTransforms, ys::Tuple)
-    index = firstindex(x)
+function _inverse!_tuple(x::AbstractVector, index, ts::NTransforms, ys::Tuple)
     for (t, y) in zip(ts, ys)
-        d = dimension(t)
-        inverse!(view_into(x, index, d), t, y)
-        index += d
+        index = inverse_at!(x, index, t, y)
     end
-    x
+    index
 end
 
-transform_with(flag::LogJacFlag, tt::TransformTuple{<:Tuple}, x::RealVector) =
-    transform_tuple(flag, tt.transformations, x)
+function transform_with(flag::LogJacFlag, tt::TransformTuple{<:Tuple}, x, index)
+    transform_tuple(flag, tt.transformations, x, index)
+end
 
 function inverse_eltype(tt::TransformTuple{<:Tuple}, y::Tuple)
     @unpack transformations = tt
     @argcheck length(transformations) == length(y)
     _inverse_eltype_tuple(transformations, y)
 end
 
-function inverse!(x::RealVector, tt::TransformTuple{<:Tuple}, y::Tuple)
+function inverse_at!(x::AbstractVector, index, tt::TransformTuple{<:Tuple}, y::Tuple)
     @unpack transformations = tt
     @argcheck length(transformations) == length(y)
     @argcheck length(x) == dimension(tt)
-    _inverse!_tuple(x, tt.transformations, y)
+    _inverse!_tuple(x, index, tt.transformations, y)
 end
 
 as(transformations::NamedTuple{N,<:NTransforms}) where N =
     TransformTuple(transformations)
 
-function transform_with(flag::LogJacFlag, tt::TransformTuple{<:NamedTuple}, x::RealVector)
+function transform_with(flag::LogJacFlag, tt::TransformTuple{<:NamedTuple}, x, index)
     @unpack transformations = tt
-    y, ℓ = transform_tuple(flag, values(transformations), x)
-    NamedTuple{keys(transformations)}(y), ℓ
+    y, ℓ, index′ = transform_tuple(flag, values(transformations), x, index)
+    NamedTuple{keys(transformations)}(y), ℓ, index′
 end
 
 function inverse_eltype(tt::TransformTuple{<:NamedTuple}, y::NamedTuple)
@@ -230,9 +227,9 @@ function inverse_eltype(tt::TransformTuple{<:NamedTuple}, y::NamedTuple)
     _inverse_eltype_tuple(values(transformations), values(y))
 end
 
-function inverse!(x::RealVector, tt::TransformTuple{<:NamedTuple}, y::NamedTuple)
+function inverse_at!(x::AbstractVector, index, tt::TransformTuple{<:NamedTuple}, y::NamedTuple)
     @unpack transformations = tt
     @argcheck keys(transformations) == keys(y)
     @argcheck length(x) == dimension(tt)
-    _inverse!_tuple(x, values(transformations), values(y))
+    _inverse!_tuple(x, index, values(transformations), values(y))
 end

src/custom.jl

@@ -74,14 +74,16 @@ CustomTransform(n::Integer, f, flatten; kwargs...) =
 
 dimension(t::CustomTransform) = dimension(t.g)
 
-function transform_with(flag::NoLogJac, t::CustomTransform, x::RealVector)
+function transform_with(flag::NoLogJac, t::CustomTransform, x::AbstractVector, index)
     @unpack g, f = t
-    f(first(transform_with(flag, g, x))), flag
+    f(first(transform_with(flag, g, x, index))), flag, index + dimension(t)
 end
 
-function transform_with(flag::LogJac, t::CustomTransform, x::RealVector)
+function transform_with(flag::LogJac, t::CustomTransform, x::AbstractVector, index)
     @unpack g, f, flatten, cfg = t
     index = firstindex(x)
-    xv = @view x[index:(index + dimension(g) - 1)]
-    value_and_logjac_forwarddiff(_custom_f(g, f), xv; flatten = flatten, cfg = cfg)
+    index′ = index + dimension(g)
+    y, ℓ = value_and_logjac_forwarddiff(_custom_f(g, f), x[index:(index′ - 1)];
+                                        flatten = flatten, cfg = cfg)
+    y, ℓ, index′
 end

src/generic.jl

@@ -1,9 +1,9 @@
 export dimension, transform, transform_and_logjac, transform_logdensity, inverse, inverse!,
     inverse_eltype, as, random_arg, random_value
 
-####
-#### log absolute Jacobian determinant
-####
+###
+### log absolute Jacobian determinant
+###
 
 """
 $(TYPEDEF)
@@ -48,8 +48,47 @@ logjac_zero(::LogJac, T::Type{<:Real}) = log(one(T))
 
 logjac_zero(::NoLogJac, _) = NOLOGJAC
 
+###
+### internal methods that implement transformations
+###
+
+"""
+    transform_with(flag::LogJacFlag, transformation, x::AbstractVector, index)
+
+Transform elements of `x` from `index`, using `transformation`.
+
+Return `(y, logjac), index′`, where
+
+- `y` is the result of the transformation,
+
+- `logjac` is the the log Jacobian determinant or a placeholder, depending on `flag`,
+
+- `index′` is the next index in `x` after the elements used for the transformation
+
+**Internal function**. Implementations
+
+1. can assume that `x` has enough elements for `transformation` (ie `@inbounds` can be
+used),
+
+2. should work with generalized indexing on `x`.
+"""
+function transform_with end
+
+"""
+    inverse_at!(x, index, transformation, y)
+
+Invert transformation at `y` and put the result in `x` starting at `index`.
+
+**Internal function**. Implementations
+
+1. can assume that `x` has enough elements for the result (ie `@inbounds` can be used),
+
+2. should work with generalized indexing on `x`.
+"""
+function inverse_at! end
+
 ####
-#### general
+#### API
 ####
 
 """
@@ -69,22 +108,6 @@ The user interface consists of
 """
 abstract type AbstractTransform end
 
-"""
-    transform_with(flag::LogJacFlag, t::AbstractTransform, x::RealVector)
-
-Transform elements of `x`, starting using `transformation`.
-
-The first value returned is the transformed value, the second the log Jacobian
-determinant or a placeholder, depending on `flag`.
-
-In contrast to [`transform`] and [`transform_and_logjac`], this method always
-assumes that `x` is a `RealVector`, for efficient traversal. Some types
-implement the latter two via this method.
-
-Implementations should assume generalized indexing on `x`.
-"""
-function transform_with end
-
 """
 $(TYPEDEF)
 
@@ -116,15 +139,19 @@ The element type for vector `x` so that `inverse!(x, t, y)` works.
 function inverse_eltype end
 
 """
-    inverse!(x, t::AbstractTransform, y)
+$(SIGNATURES)
 
 Put `inverse(t, y)` into a preallocated vector `x`, returning `x`.
 
 Generalized indexing should be assumed on `x`.
 
 See [`inverse_eltype`](@ref) for determining the type of `x`.
 """
-function inverse! end
+function inverse!(x::AbstractVector, transformation::AbstractTransform, y)
+    @argcheck dimension(transformation) == length(x)
+    inverse_at!(x, firstindex(x), transformation, y)
+    x
+end
 
 """
 $(SIGNATURES)
@@ -169,17 +196,10 @@ function as end
 #### vector transformations
 ####
 
-"""
-An `AbstractVector` of `<:Real` elements.
-
-Used internally as a type for transformations from vectors.
-"""
-const RealVector{T <: Real} = AbstractVector{T}
-
 """
 $(TYPEDEF)
 
-Transformation that transforms `<: RealVector`s to other values.
+Transformation that transforms `<: AbstractVector`s to other values.
 
 # Implementation
 
@@ -193,18 +213,26 @@ $(SIGNATURES)
 
 Transform `x` using `t`.
 """
-transform(t::VectorTransform, x::RealVector) = first(transform_with(NOLOGJAC, t, x))
+function transform(t::VectorTransform, x::AbstractVector)
+    @argcheck dimension(t) == length(x)
+    first(transform_with(NOLOGJAC, t, x, firstindex(x)))
+end
 
 """
 $(SIGNATURES)
 
 Transform `x` using `t`; calculating the log Jacobian determinant, returned as
 the second value.
 """
-transform_and_logjac(t::VectorTransform, x::RealVector) = transform_with(LOGJAC, t, x)
+function transform_and_logjac(t::VectorTransform, x::AbstractVector)
+    @argcheck dimension(t) == length(x)
+    y, ℓ, _ = transform_with(LOGJAC, t, x, firstindex(x))
+    y, ℓ
+end
 
-inverse(t::VectorTransform, y) =
+function inverse(t::VectorTransform, y)
     inverse!(Vector{inverse_eltype(t, y)}(undef, dimension(t)), t, y)
+end
 
 """
 $(SIGNATURES)

src/scalar.jl

@@ -13,14 +13,17 @@ abstract type ScalarTransform <: AbstractTransform end
 
 dimension(::ScalarTransform) = 1
 
-transform_with(flag::NoLogJac, t::ScalarTransform, x::RealVector) =
-    transform(t, @inbounds first(x)), flag
+function transform_with(flag::NoLogJac, t::ScalarTransform, x::AbstractVector, index)
+    transform(t, @inbounds x[index]), flag, index + 1
+end
 
-transform_with(::LogJac, t::ScalarTransform, x::RealVector) =
-    transform_and_logjac(t, @inbounds first(x))
+function transform_with(::LogJac, t::ScalarTransform, x::AbstractVector, index)
+    transform_and_logjac(t, @inbounds x[index])..., index + 1
+end
 
-function inverse!(x::RealVector, t::ScalarTransform, y::Real)
-    x[firstindex(x)] = inverse(t, y)
+function inverse_at!(x::AbstractVector, index, t::ScalarTransform, y::Real)
+    x[index] = inverse(t, y)
+    index + 1
 end
 
 inverse_eltype(t::ScalarTransform, y::T) where {T <: Real} = float(T)