Rewrite: removing dimensionality and allow non-bijective transformations

Project.toml

@@ -4,9 +4,11 @@ version = "0.10.0"
 
 [deps]
 ArgCheck = "dce04be8-c92d-5529-be00-80e4d2c0e197"
+Batching = "a4738bf4-2ff1-4f03-87a2-28fa6f9d5d14"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 ChangesOfVariables = "9e997f8a-9a97-42d5-a9f1-ce6bfc15e2c0"
 Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
+ConstructionBase = "187b0558-2788-49d3-abe0-74a17ed4e7c9"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
 InverseFunctions = "3587e190-3f89-42d0-90ee-14403ec27112"

docs/Project.toml

@@ -0,0 +1,3 @@
+[deps]
+Bijectors = "76274a88-744f-5084-9051-94815aaf08c4"
+Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"

docs/make.jl

@@ -0,0 +1,15 @@
+using Documenter
+using Bijectors
+
+makedocs(
+    sitename = "Bijectors",
+    format = Documenter.HTML(),
+    modules = [Bijectors]
+)
+
+# Documenter can also automatically deploy documentation to gh-pages.
+# See "Hosting Documentation" and deploydocs() in the Documenter manual
+# for more information.
+#=deploydocs(
+    repo = "<repository url>"
+)=#

docs/src/index.md

@@ -0,0 +1,144 @@
+# Bijectors.jl
+
+Documentation for Bijectors.jl
+
+## Usage
+
+A very simple example of a "bijector"/diffeomorphism, i.e. a differentiable transformation with a differentiable inverse, is the `exp` function:
+- The inverse of `exp` is `log`.
+- The derivative of `exp` at an input `x` is simply `exp(x)`, hence `logabsdetjac` is simply `x`.
+
+```@repl usage
+using Bijectors
+transform(exp, 1.0)
+logabsdetjac(exp, 1.0)
+forward(exp, 1.0)
+```
+
+If you want to instead transform a collection of inputs, you can use the `batch` method from Batching.jl to inform Bijectors.jl that the input now represents a collection of inputs rather than a single input:
+
+```@repl usage
+xs = batch(ones(2));
+transform(exp, xs)
+logabsdetjac(exp, xs)
+forward(exp, xs)
+```
+
+Some transformations is well-defined for different types of inputs, e.g. `exp` can also act elementwise on a `N`-dimensional `Array{<:Real,N}`. To specify that a transformation should be acting elementwise, we use the [`elementwise`](@ref) method:
+
+```@repl usage
+x = ones(2, 2)
+transform(elementwise(exp), x)
+logabsdetjac(elementwise(exp), x)
+forward(elementwise(exp), x)
+```
+
+And batched versions:
+
+```@repl usage
+xs = batch(ones(2, 2, 3));
+transform(elementwise(exp), xs)
+logabsdetjac(elementwise(exp), xs)
+forward(elementwise(exp), xs)
+```
+
+These methods also work nicely for compositions of transformations:
+
+```@repl usage
+transform(elementwise(log ∘ exp), xs)
+```
+
+Unlike `exp`, some transformations have parameters affecting the resulting transformation they represent, e.g. `Logit` has two parameters `a` and `b` representing the lower- and upper-bound, respectively, of its domain:
+
+```@repl usage
+using Bijectors: Logit
+
+f = Logit(0.0, 1.0)
+f(rand()) # takes us from `(0, 1)` to `(-∞, ∞)`
+```
+
+## User-facing methods
+
+Without mutation:
+
+```@docs
+transform
+logabsdetjac
+forward(b, x)
+```
+
+With mutation:
+
+```@docs
+transform!
+logabsdetjac!
+forward!
+```
+
+## Implementing a transformation
+
+Any callable can be made into a bijector by providing an implementation of [`forward(b, x)`](@ref), which is done by overloading
+
+```@docs
+Bijectors.forward_single
+```
+
+where
+
+```@docs
+Bijectors.transform_single
+Bijectors.logabsdetjac_single
+```
+
+You can then optionally implement `transform` and `logabsdetjac` to avoid redundant computations. This is usually only worth it if you expect `transform` or `logabsdetjac` to be used heavily without the other.
+
+Note that a _user_ of the bijector should generally be using [`forward(b, x)`](@ref) rather than calling [`forward_single`](@ref) directly.
+
+To implement "batched" versions of the above functionalities, i.e. methods which act on a _collection_ of inputs rather than a single input, you can overload the following method:
+
+```@docs
+Bijectors.forward_multiple
+```
+
+And similarly, if you want to specialize [`transform`](@ref) and [`logabsdetjac`](@ref), you can implement
+
+```@docs
+Bijectors.transform_multiple
+Bijectors.logabsdetjac_multiple
+```
+
+### Mutability
+
+There are also _mutable_ versions of all of the above:
+
+```@docs
+Bijectors.forward_single!
+Bijectors.forward_multiple!
+Bijectors.transform_single!
+Bijectors.transform_multiple!
+Bijectors.logabsdetjac_single!
+Bijectors.logabsdetjac_multiple!
+```
+
+## Working with Distributions.jl
+
+```@docs
+Bijectors.bijector
+Bijectors.transformed(d::Distribution, b::Bijector)
+```
+
+## Utilities
+
+```@docs
+Bijectors.elementwise
+Bijectors.isinvertible
+Bijectors.isclosedform(t::Bijectors.Transform)
+```
+
+## API
+
+```@docs
+Bijectors.Transform
+Bijectors.Bijector
+Bijectors.Inverse
+```

src/Bijectors.jl

@@ -32,9 +32,13 @@ using Reexport, Requires
 @reexport using Distributions
 using LinearAlgebra
 using MappedArrays
+using ConstructionBase
 using Base.Iterators: drop
 using LinearAlgebra: AbstractTriangular
 
+using Batching
+export batch
+
 import ChangesOfVariables: with_logabsdet_jacobian
 import InverseFunctions: inverse
 
@@ -54,10 +58,13 @@ export  TransformDistribution,
         logpdf_with_trans,
         isclosedform,
         transform,
+        transform!,
         with_logabsdet_jacobian,
         inverse,
         forward,
+        forward!,
         logabsdetjac,
+        logabsdetjac!,
         logabsdetjacinv,
         Bijector,
         ADBijector,
@@ -76,7 +83,9 @@ export  TransformDistribution,
         PlanarLayer,
         RadialLayer,
         CouplingLayer,
-        InvertibleBatchNorm
+        InvertibleBatchNorm,
+        Elementwise,
+        elementwise
 
 if VERSION < v"1.1"
     using Compat: eachcol
@@ -251,12 +260,13 @@ function getlogp(d::InverseWishart, Xcf, X)
 end
 
 include("utils.jl")
+# include("batch.jl")
 include("interface.jl")
-include("chainrules.jl")
+# include("chainrules.jl")
 
-Base.@deprecate forward(b::AbstractBijector, x) NamedTuple{(:rv,:logabsdetjac)}(with_logabsdet_jacobian(b, x))
+Base.@deprecate forward(b::Transform, x) NamedTuple{(:rv,:logabsdetjac)}(with_logabsdet_jacobian(b, x))
 
-@noinline function Base.inv(b::AbstractBijector)
+@noinline function Base.inv(b::Transform)
     Base.depwarn("`Base.inv(b::AbstractBijector)` is deprecated, use `inverse(b)` instead.", :inv)
     inverse(b)
 end
@@ -265,24 +275,24 @@ end
 maporbroadcast(f, x::AbstractArray{<:Any, N}...) where {N} = map(f, x...)
 maporbroadcast(f, x::AbstractArray...) = f.(x...)
 
-# optional dependencies
-function __init__()
-    @require LazyArrays = "5078a376-72f3-5289-bfd5-ec5146d43c02" begin
-        function maporbroadcast(f, x1::LazyArrays.BroadcastArray, x...)
-            return copy(f.(x1, x...))
-        end
-        function maporbroadcast(f, x1, x2::LazyArrays.BroadcastArray, x...)
-            return copy(f.(x1, x2, x...))
-        end
-        function maporbroadcast(f, x1, x2, x3::LazyArrays.BroadcastArray, x...)
-            return copy(f.(x1, x2, x3, x...))
-        end
-    end
-    @require ForwardDiff="f6369f11-7733-5829-9624-2563aa707210" include("compat/forwarddiff.jl")
-    @require Tracker="9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c" include("compat/tracker.jl")
-    @require Zygote="e88e6eb3-aa80-5325-afca-941959d7151f" include("compat/zygote.jl")
-    @require ReverseDiff="37e2e3b7-166d-5795-8a7a-e32c996b4267" include("compat/reversediff.jl")
-    @require DistributionsAD="ced4e74d-a319-5a8a-b0ac-84af2272839c" include("compat/distributionsad.jl")
-end
+# # optional dependencies
+# function __init__()
+#     @require LazyArrays = "5078a376-72f3-5289-bfd5-ec5146d43c02" begin
+#         function maporbroadcast(f, x1::LazyArrays.BroadcastArray, x...)
+#             return copy(f.(x1, x...))
+#         end
+#         function maporbroadcast(f, x1, x2::LazyArrays.BroadcastArray, x...)
+#             return copy(f.(x1, x2, x...))
+#         end
+#         function maporbroadcast(f, x1, x2, x3::LazyArrays.BroadcastArray, x...)
+#             return copy(f.(x1, x2, x3, x...))
+#         end
+#     end
+#     @require ForwardDiff="f6369f11-7733-5829-9624-2563aa707210" include("compat/forwarddiff.jl")
+#     @require Tracker="9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c" include("compat/tracker.jl")
+#     @require Zygote="e88e6eb3-aa80-5325-afca-941959d7151f" include("compat/zygote.jl")
+#     @require ReverseDiff="37e2e3b7-166d-5795-8a7a-e32c996b4267" include("compat/reversediff.jl")
+#     @require DistributionsAD="ced4e74d-a319-5a8a-b0ac-84af2272839c" include("compat/distributionsad.jl")
+# end
 
 end # module

src/batch.jl

@@ -0,0 +1,88 @@
+abstract type AbstractBatch{T} <: AbstractVector{T} end
+
+"""
+    value(x)
+
+Returns the underlying storage used for the entire batch.
+
+If `x` is not `AbstractBatch`, then this is the identity function.
+"""
+value(x) = x
+
+struct Batch{V, T} <: AbstractBatch{T}
+    value::V
+end
+
+value(x::Batch) = x.value
+
+# Convenient aliases
+const ArrayBatch{N} = Batch{<:AbstractArray{<:Real, N}}
+const VectorBatch{T} = Batch{<:AbstractVector{T}}
+
+# Constructor for `ArrayBatch`.
+Batch(x::AbstractVector{<:Real}) = Batch{typeof(x), eltype(x)}(x)
+function Batch(x::AbstractArray{<:Real})
+    V = typeof(x)
+    # HACK: This assumes the batch is non-empty.
+    T = typeof(getindex_for_last(x, 1))
+    return Batch{V, T}(x)
+end
+
+# Constructor for `VectorBatch`.
+Batch(x::AbstractVector{<:AbstractArray}) = Batch{typeof(x), eltype(x)}(x)
+
+# `AbstractVector` interface.
+Base.size(batch::Batch{<:AbstractArray{<:Any, N}}) where {N} = (size(value(batch), N), )
+
+# Since impl inherited from `AbstractVector` doesn't do exactly what we want.
+Base.similar(b::AbstractBatch) = reconstruct(b, similar(value(b)))
+
+# For `VectorBatch`
+Base.getindex(batch::VectorBatch, i::Int) = value(batch)[i]
+Base.getindex(batch::VectorBatch, i::CartesianIndex{1}) = value(batch)[i]
+Base.getindex(batch::VectorBatch, i) = Batch(value(batch)[i])
+function Base.setindex!(batch::VectorBatch, v, i)
+    # `v` can also be a `Batch`.
+    Base.setindex!(value(batch), value(v), i)
+    return batch
+end
+
+# For `ArrayBatch`
+@generated function getindex_for_last(x::AbstractArray{<:Any, N}, inds) where {N}
+    e = Expr(:call)
+    push!(e.args, :(Base.view))
+    push!(e.args, :x)
+
+    for i = 1:N - 1
+        push!(e.args, :(:))
+    end
+
+    push!(e.args, :(inds))
+
+    return e
+end
+
+@generated function setindex_for_last!(out::AbstractArray{<:Any, N}, x, inds) where {N}
+    e = Expr(:call)
+    push!(e.args, :(Base.setindex!))
+    push!(e.args, :out)
+    push!(e.args, :x)
+
+    for i = 1:N - 1
+        push!(e.args, :(:))
+    end
+
+    push!(e.args, :(inds))
+
+    return e
+end
+
+# General arrays.
+Base.getindex(batch::ArrayBatch, i::Int) = getindex_for_last(value(batch), i)
+Base.getindex(batch::ArrayBatch, i::CartesianIndex{1}) = getindex_for_last(value(batch), i)
+Base.getindex(batch::ArrayBatch, i) = Batch(getindex_for_last(value(batch), i))
+function Base.setindex!(batch::ArrayBatch, v, i)
+    # `v` can also be a `Batch`.
+    setindex_for_last!(value(batch), value(v), i)
+    return batch
+end

src/bijectors/adbijector.jl

@@ -2,7 +2,7 @@
 Abstract type for a `Bijector{N}` making use of auto-differentation (AD) to
 implement `jacobian` and, by impliciation, `logabsdetjac`.
 """
-abstract type ADBijector{AD, N} <: Bijector{N} end
+abstract type ADBijector{AD} <: Bijector end
 
 struct SingularJacobianException{B<:Bijector} <: Exception
     b::B
@@ -24,4 +24,4 @@ end
 function logabsdetjac(b::ADBijector, x::AbstractVector{<:Real})
     fact = lu(jacobian(b, x), check=false)
     return issuccess(fact) ? logabsdet(fact)[1] : throw(SingularJacobianException(b))
-end
+end