Support Broadcast.BroadcastFunction

Author: oschulz

docs/src/index.md

@@ -6,6 +6,6 @@ ChangesOfVariables
 
 This package defines the function [`with_logabsdet_jacobian`](@ref). `(y, ladj) = with_logabsdet_jacobian(f, x)` computes both the transformed value of `x` under the transformation `f` and the logarithm of the [volume element](https://en.wikipedia.org/wiki/Volume_element).
 
-`with_logabsdet_jacobian` supports mapped/broadcasted functions (via `Base.Fix1`) and (on Julia >=v1.6) function composition.
+`with_logabsdet_jacobian` supports mapped/broadcasted functions (via `Base.Broadcast.BroadcastFunction` or `Base.Fix1`) and (on Julia >=v1.6) function composition.
 
 Implementations of `with_logabsdet_jacobian(f)` for `identity`, `inv`, `adjoint` and `transpose` as well as for `exp`, `log`, `exp2`, `log2`, `exp10`, `log10`, `expm1` and `log1p` are included.

src/with_ladj.jl

@@ -12,7 +12,8 @@ For `(y, ladj) = with_logabsdet_jacobian(f, x)`, the following must hold true:
 * `ladj` is the `log(abs(det(jacobian(f, x))))`
 
 `with_logabsdet_jacobian` comes with support for broadcasted/mapped functions
-(via `Base.Fix1`) and (Julia >=v1.6 only) `ComposedFunction`.
+(via `Base.Broadcast.BroadcastFunction` or `Base.Fix1`) and (Julia >=v1.6 only)
+`ComposedFunction`.
 
 If no volume element is defined/applicable, `with_logabsdet_jacobian(f::F, x::T)`
 returns [`NoLogAbsDetJacobian{F,T}()`](@ref).
@@ -43,7 +44,11 @@ true
 
 ```jldoctest a
 X = rand(10)
-broadcasted_foo = Base.Fix1(broadcast, foo)
+broadcasted_foo = if VERSION >= v"1.6"
+    Base.Broadcast.BroadcastFunction(foo)
+else
+    Base.Fix1(broadcast, foo)
+end
 Y, ladj_Y = with_logabsdet_jacobian(broadcasted_foo, X)
 Y == broadcasted_foo(X) && ladj_Y ≈ logabsdet(ForwardDiff.jacobian(broadcasted_foo, X))[1]
 
@@ -117,6 +122,13 @@ function _with_ladj_on_mapped(map_or_bc::F, y_with_ladj) where {F<:Union{typeof(
     (y, ladj)
 end
 
+@static if VERSION >= v"1.6"
+    function with_logabsdet_jacobian(mapped_f::Base.Broadcast.BroadcastFunction, X)
+        f = mapped_f.f
+        y_with_ladj = broadcast(Base.Fix1(with_logabsdet_jacobian, f), X)
+        _with_ladj_on_mapped(broadcast, y_with_ladj)
+    end
+end
 
 function with_logabsdet_jacobian(mapped_f::Base.Fix1{<:Union{typeof(map),typeof(broadcast)}}, X)
     map_or_bc = mapped_f.f

test/test_with_ladj.jl

@@ -12,6 +12,12 @@ include("getjacobian.jl")
 
 
 @testset "with_logabsdet_jacobian" begin
+    @static if VERSION >= v"1.6"
+        _bc_func(f) = Base.Broadcast.BroadcastFunction(f)
+    else
+        _bc_func(f) = Base.Fix1(broadcast, f)
+    end
+
     @test with_logabsdet_jacobian(sum, rand(5)) == NoLogAbsDetJacobian{typeof(sum),Vector{Float64}}()
     @test with_logabsdet_jacobian(sum ∘ log, 5.0f0) == NoLogAbsDetJacobian{typeof(sum ∘ log),Float32}()
     @test with_logabsdet_jacobian(log ∘ sum, 5.0f0) == NoLogAbsDetJacobian{typeof(log ∘ sum),Float32}()
@@ -39,7 +45,7 @@ include("getjacobian.jl")
     end
 
     @testset "with_logabsdet_jacobian on mapped and broadcasted" begin
-        for f in (Base.Fix1(map, foo), Base.Fix1(broadcast, foo))
+        for f in (_bc_func(foo), Base.Fix1(map, foo), Base.Fix1(broadcast, foo))
             for arg in (x, fill(x,), Ref(x), (x,), X)
                 test_with_logabsdet_jacobian(f, arg, getjacobian, compare = isaprx)
             end