Add support for distributions with monotonically increasing bijector

src/bijectors/exp_log.jl

@@ -7,3 +7,6 @@ logabsdetjac(b::Elementwise{typeof(exp)}, x) = sum(x)
 
 logabsdetjac(b::typeof(log), x::Real) = -log(x)
 logabsdetjac(b::Elementwise{typeof(log)}, x) = -sum(log, x)
+
+is_monotonically_increasing(::typeof(exp)) = true
+is_monotonically_increasing(::typeof(log)) = true

src/bijectors/leaky_relu.jl

@@ -27,3 +27,5 @@ function with_logabsdet_jacobian(b::LeakyReLU, x::AbstractArray)
     J = mask .* b.α .+ (!).(mask)
     return J .* x, sum(log.(abs.(J)))
 end
+
+is_monotonically_increasing(::LeakyReLU) = true

src/bijectors/logit.jl

@@ -28,3 +28,6 @@ logabsdetjac(b::Logit, x) = sum(logit_logabsdetjac.(x, b.a, b.b))
 function with_logabsdet_jacobian(b::Logit, x)
     return _logit.(x, b.a, b.b), sum(logit_logabsdetjac.(x, b.a, b.b))
 end
+
+
+is_monotonically_increasing(::Logit) = true

src/bijectors/ordered.jl

@@ -1,3 +1,11 @@
+struct SignFlip <: Bijector end
+
+with_logabsdet_jacobian(::SignFlip, x) = -x, zero(eltype(x))
+inverse(::SignFlip) = SignFlip()
+output_size(::SignFlip, x) = size(x)
+is_monotonically_increasing(::SignFlip) = false
+is_monotonically_decreasing(::SignFlip) = false
+
 """
     OrderedBijector()
 
@@ -17,14 +25,22 @@ 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 bijector(d) !== identity
-        throw(
-            ArgumentError(
-                "ordered transform is currently only supported for unconstrained distributions.",
-            ),
-        )
+    # We're good if the map from unconstrained (in which we apply the ordered bijector)
+    # to constrained is monotonically increasing, i.e. order-preserving. In that case,
+    # we can form the ordered transformation as `binv ∘ OrderedBijector() ∘ b`.
+    # TODO: Add support for monotonically _decreasing_ transformations. This will be the
+    # the same as above, but the ordering will be reversed by `binv` so we need to handle this.
+    b = bijector(d)
+    binv = inverse(b)
+    if is_monotonically_decreasing(binv)
+        ordered_b = binv ∘ SignFlip() ∘ OrderedBijector() ∘ SignFlip() ∘ b
+    elseif is_monotonically_increasing(binv)
+        ordered_b = binv ∘ OrderedBijector() ∘ b
+    else
+        throw(ArgumentError("ordered transform is currently not supported for $d."))
     end
-    return transformed(d, OrderedBijector())
+
+    return transformed(d, ordered_b)
 end
 
 with_logabsdet_jacobian(b::OrderedBijector, x) = transform(b, x), logabsdetjac(b, x)

src/bijectors/scale.jl

@@ -34,3 +34,6 @@ _logabsdetjac_scale(a::AbstractVector, x::AbstractMatrix, ::Val{2}) = sum(log 
 # Matrix: single input.
 _logabsdetjac_scale(a::AbstractMatrix, x::AbstractVector, ::Val{1}) = logabsdet(a)[1]
 _logabsdetjac_scale(a::AbstractMatrix, x::AbstractMatrix, ::Val{2}) = logabsdet(a)[1]
+
+is_monotonically_increasing(a::Scale) = all(Base.Fix1(>, 0), a.a)
+is_monotonically_decreasing(a::Scale) = all(Base.Fix1(<, 0), a.a)

src/bijectors/shift.jl

@@ -22,3 +22,5 @@ _logabsdetjac_shift(a, x) = zero(eltype(x))
 _logabsdetjac_shift_array_batch(a, x) = zeros(eltype(x), size(x, ndims(x)))
 
 with_logabsdet_jacobian(b::Shift, x) = transform(b, x), logabsdetjac(b, x)
+
+is_monotonically_increasing(::Shift) = true

src/bijectors/truncated.jl

@@ -67,3 +67,36 @@ function truncated_logabsdetjac(x, a, b)
 end
 
 with_logabsdet_jacobian(b::TruncatedBijector, x) = transform(b, x), logabsdetjac(b, x)
+
+# It's only monotonically decreasing if it's only upper-bounded.
+# In the multivariate case, we can only say something reasonable if entries are monotonic.
+function is_monotonically_increasing(b::TruncatedBijector)
+    lowerbounded, upperbounded = all(isfinite, b.lb), all(isfinite, b.ub)
+    return if lowerbounded
+        true
+    elseif upperbounded
+        # => decreasing
+        false
+    elseif all(!isfinite, b.lb) && all(!isfinite, b.ub)
+        # => all are unbounded so we have the identity
+        true
+    else
+        # => some are unbounded and some are bounded
+        false
+    end
+end
+function is_monotonically_decreasing(b::TruncatedBijector)
+    lowerbounded, upperbounded = all(isfinite, b.lb), all(isfinite, b.ub)
+    return if lowerbounded
+        false
+    elseif upperbounded
+        # => decreasing
+        true
+    elseif all(!isfinite, b.lb) && all(!isfinite, b.ub)
+        # => all are unbounded so we have the identity
+        false
+    else
+        # => some are unbounded and some are bounded
+        true
+    end
+end

src/interface.jl

@@ -231,6 +231,63 @@ transform!(::typeof(identity), x, y) = copy!(y, x)
 logabsdetjac(::typeof(identity), x) = zero(eltype(x))
 logabsdetjac!(::typeof(identity), x, logjac) = logjac
 
+###################
+# Other utilities #
+###################
+"""
+    is_monotonically_increasing(f)
+
+Returns `true` if `f` is monotonically increasing.
+"""
+is_monotonically_increasing(f) = false
+is_monotonically_increasing(::typeof(identity)) = true
+is_monotonically_increasing(binv::Inverse) = is_monotonically_increasing(inverse(binv))
+is_monotonically_increasing(ef::Elementwise) = is_monotonically_increasing(ef.x)
+function is_monotonically_increasing(cf::ComposedFunction)
+    # Here we have a few different cases:
+    #
+    #     inner \ outer | inc | dec | other
+    #     --------------+-----+-----+------
+    #     inc           | inc | dec | NA   
+    #     dec           | dec | inc | NA   
+    #     other         | NA  | NA  | NA   
+    #     --------------+-----+-----+------
+    return if is_monotonically_increasing(cf.inner)
+        is_monotonically_increasing(cf.outer)
+    elseif is_monotonically_decreasing(cf.inner)
+        is_monotonically_decreasing(cf.outer)
+    else
+        false
+    end
+end
+
+"""
+    is_monotonically_decreasing(f)
+
+Returns `true` if `f` is monotonically decreasing.
+    """
+is_monotonically_decreasing(f) = false
+is_monotonically_decreasing(::typeof(identity)) = false
+is_monotonically_decreasing(binv::Inverse) = is_monotonically_decreasing(inverse(binv))
+is_monotonically_decreasing(ef::Elementwise) = is_monotonically_decreasing(ef.x)
+function is_monotonically_decreasing(cf::ComposedFunction)
+    # Here we have a few different cases:
+    #
+    #     inner \ outer | inc | dec | other
+    #     --------------+-----+-----+------
+    #     inc           | inc | dec | NA   
+    #     dec           | dec | inc | NA   
+    #     other         | NA  | NA  | NA   
+    #     --------------+-----+-----+------
+    return if is_monotonically_increasing(cf.inner)
+        is_monotonically_decreasing(cf.outer)
+    elseif is_monotonically_decreasing(cf.inner)
+        is_monotonically_increasing(cf.outer)
+    else
+        false
+    end
+end
+
 ######################
 # Bijectors includes #
 ######################

test/bijectors/ordered.jl

@@ -22,11 +22,20 @@ end
         MvTDist(1, collect(1.0:5), Matrix(I(5))),
         product_distribution(fill(Normal(), 5)),
         product_distribution(fill(TDist(1), 5)),
+        # positive supports
+        product_distribution(fill(LogNormal(), 5)),
+        product_distribution(fill(InverseGamma(2, 3), 5)),
+        # negative supports
+        product_distribution(fill(-1 * InverseGamma(2, 3), 5)),
+        # bounded supports
+        product_distribution(fill(Uniform(1, 2), 5)),
+        # different transformations
+        product_distribution(fill(transformed(InverseGamma(2, 3), Bijectors.Scale(3)), 5)),
+        product_distribution(fill(transformed(InverseGamma(2, 3), Bijectors.Scale(-3)), 5)),
     ]
         d_ordered = ordered(d)
         @test d_ordered isa Bijectors.TransformedDistribution
         @test d_ordered.dist === d
-        @test d_ordered.transform isa OrderedBijector
         y = randn(5)
         x = inverse(bijector(d_ordered))(y)
         @test issorted(x)