Loosen types for Reactant

.github/workflows/CI.yml

@@ -17,7 +17,7 @@ jobs:
       fail-fast: false
       matrix:
         version:
-          - 'min'
+          - 'lts'
           - '1'
         os:
           - ubuntu-latest
@@ -54,7 +54,7 @@ jobs:
       fail-fast: false
       matrix:
         version:
-          - 'min'
+          - 'lts'
           - '1'
         os:
           - ubuntu-latest

.gitignore

@@ -7,3 +7,4 @@
 /docs/Manifest.toml
 /test/coverage/Manifest.toml
 /test/Manifest.toml
+.vscode/settings.json

Project.toml

@@ -18,11 +18,13 @@ StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697"
 ChangesOfVariables = "9e997f8a-9a97-42d5-a9f1-ce6bfc15e2c0"
 InverseFunctions = "3587e190-3f89-42d0-90ee-14403ec27112"
+Reactant = "3c362404-f566-11ee-1572-e11a4b42c853"
 
 [extensions]
 AccessorsExt = "Accessors"
 ChangesOfVariablesExt = "ChangesOfVariables"
 InverseFunctionsExt = "InverseFunctions"
+ReactantExt = "Reactant"
 
 [compat]
 Accessors = "0.1.42"
@@ -36,6 +38,7 @@ InverseFunctions = "0.1"
 LinearAlgebra = "1.6"
 LogExpFunctions = "0.3"
 Random = "1.6"
+Reactant = "0.2.229"
 StaticArrays = "1"
 julia = "1.10"
 

ext/ReactantExt.jl

@@ -0,0 +1,92 @@
+module ReactantExt
+using TransformVariables: TransformVariables, ArrayTransformation, LogJacFlag,  
+                          logjac_zero, transform_with, _ensure_float, dimension
+using Reactant
+using Reactant: TracedRNumber, AnyTracedRArray
+
+RInt = Union{Int, TracedRNumber{Int}}
+Base.@propagate_inbounds function TransformVariables.tv_getindex(a::AnyTracedRArray, i::RInt)
+    Reactant.@allowscalar a[i]
+end
+
+# The dims of the ArrayTransformation must be constant because Reactant can only deal with non-dynamic arrays.
+Base.@nospecialize function Reactant.traced_type_inner(
+    @nospecialize(prev::Type{ArrayTransformation{T, M}}),
+    seen,
+    @nospecialize(mode::TraceMode),
+    @nospecialize(track_numbers::Type),
+    @nospecialize(sharding),
+    @nospecialize(runtime)
+) where {T, M}
+    T_traced = Reactant.traced_type_inner(T, seen, mode, track_numbers, sharding, runtime)
+    return TransformVariables.ArrayTransformation{T_traced, M}
+end
+
+TransformVariables._ensure_float(x::Type{T}) where {T<:TracedRNumber} = T
+
+@noinline function TransformVariables.transform_with(
+    flag::TransformVariables.LogJacFlag, 
+    transformation::TransformVariables.ArrayTransformation, 
+    x::AnyTracedRArray, 
+    index
+)
+    (; inner_transformation, dims) = transformation
+    # NOTE not using index increments as that somehow breaks type inference
+    d = dimension(inner_transformation) # length of an element transformation
+    len = prod(dims)              # number of elements
+    𝐼 = range(index; length = len, step = d)
+
+    # # Reactant can't easily handle the usual return type because the array eltype is not a Reacant primitive
+    # # so we have to do a 2-pass algorithm
+    # yℓ = map(index -> ((y, ℓ, _) = transform_with(flag, inner_transformation, x, index); y), 𝐼)
+    # ℓa = map(index -> ((y, ℓ, _) = transform_with(flag, inner_transformation, x, index); ℓ), 𝐼)
+
+    # index′ = index + d * len
+    # y = reshape(yℓ, dims)
+    # ℓa = sum(ℓa)
+    # y, ℓa, index′
+
+    tmp,_,_ = transform_with(flag, inner_transformation, x, first(𝐼))
+    ℓa = logjac_zero(flag, _ensure_float(eltype(x)))
+    if typeof(tmp) <: Number
+        yℓ = similar(x, typeof(tmp), length(𝐼))
+    elseif typeof(tmp) <: AbstractArray
+        # convert to a larger array since reactant can't easily index into an julia array of Reactant arrays with a traced number.
+        yℓ = similar(tmp, size(tmp)..., length(𝐼))
+    else
+        throw(ArgumentError("Number and AbstractArray transformations are only supported in Reactant compilation mode"))
+    end
+    @trace track_numbers=false for i in eachindex(𝐼)
+        idx = 𝐼[i]
+        y, ℓ, _ = transform_with(flag, inner_transformation, x, idx)
+        if !isempty(y)
+            ℓa += ℓ        
+        end
+        if tmp isa Number
+            Reactant.@allowscalar yℓ[i] = y
+        else
+            idxs = ntuple(n-> Colon(), ndims(tmp)) 
+            yℓ[idxs..., i] = y
+        end
+    end
+
+    if tmp isa AbstractArray
+        yℓ0 = collect(eachslice(yℓ, dims=ndims(yℓ)))
+    else
+        yℓ0 = yℓ
+    end
+
+    index′ = index + d * len
+    reshape(yℓ0, dims), ℓa, index′
+end
+
+# To prevent ambiguities with regular method
+function TransformVariables.transform_with(flag::LogJacFlag, t::ArrayTransformation{TransformVariables.Identity}, x::AnyTracedRArray, index)
+    index′ = index+dimension(t)
+    inds = Reactant.TracedUnitRange(index, index′-1, dimension(t))
+    y = reshape(x[inds], t.dims)
+    y, logjac_zero(flag, _ensure_float(eltype(x))), index′
+end
+
+
+end

src/aggregation.jl

@@ -9,7 +9,7 @@ $(TYPEDEF)
 
 Apply `transformation` repeatedly to create an array with given `dims`.
 """
-struct ArrayTransformation{T <: AbstractTransform,M} <: VectorTransform
+struct ArrayTransformation{T <: AbstractTransform, M} <: VectorTransform
     inner_transformation::T
     dims::NTuple{M, Int}
 end
@@ -82,7 +82,7 @@ function transform_with(flag::LogJacFlag, transformation::ArrayTransformation, x
     first.(yℓ), isempty(yℓ) ? ℓz : ℓz + sum(last, yℓ), index′
 end
 
-function transform_with(flag::LogJacFlag, t::ArrayTransformation{Identity}, x, index)
+function transform_with(flag::LogJacFlag, t::ArrayTransformation{Identity}, x::AbstractArray, index)
     index′ = index+dimension(t)
     y = reshape(x[index:(index′-1)], t.dims)
     y, logjac_zero(flag, _ensure_float(eltype(x))), index′

src/generic.jl

@@ -50,7 +50,7 @@ $(SIGNATURES)
 
 Initial value for log Jacobian calculations.
 """
-logjac_zero(::LogJac, ::Type{T}) where {T<:Real} = log(one(T))
+logjac_zero(::LogJac, ::Type{T}) where {T<:Number} = log(one(T))
 
 logjac_zero(::NoLogJac, _) = NOLOGJAC
 

src/scalar.jl

@@ -17,12 +17,12 @@ abstract type ScalarTransform <: AbstractTransform end
 
 dimension(::ScalarTransform) = 1
 
-function transform_with(flag::NoLogJac, t::ScalarTransform, x::AbstractVector, index::Int)
-    transform(t, @inbounds x[index]), flag, index + 1
+function transform_with(flag::NoLogJac, t::ScalarTransform, x::AbstractVector, index)
+    transform(t, @inbounds tv_getindex(x, index)), flag, index + 1
 end
 
-function transform_with(::LogJac, t::ScalarTransform, x::AbstractVector, index::Int)
-    transform_and_logjac(t, @inbounds x[index])..., index + 1
+function transform_with(::LogJac, t::ScalarTransform, x::AbstractVector, index)
+    transform_and_logjac(t, @inbounds tv_getindex(x, index))..., index + 1
 end
 
 function inverse_at!(x::AbstractVector, index::Int, t::ScalarTransform, y)
@@ -43,15 +43,15 @@ Identity ``x ↦ x``.
 """
 struct Identity <: ScalarTransform end
 
-transform(::Identity, x::Real) = x
+transform(::Identity, x::Number) = x
 
-transform_and_logjac(::Identity, x::Real) = x, logjac_zero(LogJac(), typeof(x))
+transform_and_logjac(::Identity, x::Number) = x, logjac_zero(LogJac(), typeof(x))
 
 inverse_eltype(t::Identity, ::Type{T}) where T = T
 
 inverse(::Identity, x::Number) = x
 
-inverse_and_logjac(::Identity, x::Real) = x, logjac_zero(LogJac(), typeof(x))
+inverse_and_logjac(::Identity, x::Number) = x, logjac_zero(LogJac(), typeof(x))
 
 ####
 #### elementary scalar transforms
@@ -64,9 +64,9 @@ Exponential transformation `x ↦ eˣ`. Maps from all reals to the positive real
 """
 struct TVExp <: ScalarTransform end
 
-transform(::TVExp, x::Real) = exp(x)
+transform(::TVExp, x::Number) = exp(x)
 
-transform_and_logjac(t::TVExp, x::Real) = transform(t, x), x
+transform_and_logjac(t::TVExp, x::Number) = transform(t, x), x
 
 inverse_eltype(t::TVExp, ::Type{T}) where T = _ensure_float(T)
 
@@ -83,9 +83,9 @@ Logistic transformation `x ↦ logit(x)`. Maps from all reals to (0, 1).
 """
 struct TVLogistic <: ScalarTransform end
 
-transform(::TVLogistic, x::Real) = logistic(x)
+transform(::TVLogistic, x::Number) = logistic(x)
 
-transform_and_logjac(t::TVLogistic, x::Real) = transform(t, x), logistic_logjac(x)
+transform_and_logjac(t::TVLogistic, x::Number) = transform(t, x), logistic_logjac(x)
 
 inverse_eltype(t::TVLogistic, ::Type{T}) where T = _ensure_float(T)
 
@@ -100,13 +100,13 @@ $(TYPEDEF)
 
 Shift transformation `x ↦ x + shift`.
 """
-struct TVShift{T <: Real} <: ScalarTransform
+struct TVShift{T} <: ScalarTransform
     shift::T
 end
 
-transform(t::TVShift, x::Real) = x + t.shift
+transform(t::TVShift, x::Number) = x + t.shift
 
-transform_and_logjac(t::TVShift, x::Real) = transform(t, x), logjac_zero(LogJac(), typeof(x))
+transform_and_logjac(t::TVShift, x::Number) = transform(t, x), logjac_zero(LogJac(), typeof(x))
 
 inverse_eltype(t::TVShift{S}, ::Type{T}) where {S,T} = typeof(zero(_ensure_float(T)) - zero(S))
 
@@ -129,15 +129,15 @@ end
 
 TVScale(scale::T) where {T} = TVScale{T}(scale)
 
-transform(t::TVScale, x::Real) = t.scale * x
+transform(t::TVScale, x::Number) = t.scale * x
 
-transform_and_logjac(t::TVScale{<:Real}, x::Real) = transform(t, x), log(t.scale)
+transform_and_logjac(t::TVScale{<:Real}, x::Number) = transform(t, x), log(t.scale)
 
 inverse_eltype(t::TVScale{S}, ::Type{T}) where {S,T} = typeof(oneunit(T) / oneunit(S))
 
 inverse(t::TVScale, x::Number) = x / t.scale
 
-inverse_and_logjac(t::TVScale{<:Real}, x::Number) = inverse(t, x), -log(t.scale)
+inverse_and_logjac(t::TVScale, x::Number) = inverse(t, x), -log(t.scale)
 
 """
 $(TYPEDEF)
@@ -147,8 +147,8 @@ Negative transformation `x ↦ -x`.
 struct TVNeg <: ScalarTransform
 end
 
-transform(::TVNeg, x::Real) = -x
-transform_and_logjac(t::TVNeg, x::Real) = transform(t, x), logjac_zero(LogJac(), typeof(x))
+transform(::TVNeg, x::Number) = -x
+transform_and_logjac(t::TVNeg, x::Number) = transform(t, x), logjac_zero(LogJac(), typeof(x))
 
 inverse_eltype(::TVNeg, ::Type{T}) where T = typeof(-oneunit(T))
 inverse(::TVNeg, x::Number) = -x

src/special_arrays.jl

@@ -84,7 +84,7 @@ function transform_with(flag::LogJacFlag, t::UnitVector, x::AbstractVector, inde
     (; n) = t
     T = _ensure_float(eltype(x))
     log_r = zero(T)
-    y = Vector{T}(undef, n)
+    y = similar(x, T, n)
     ℓ = logjac_zero(flag, T)
     @inbounds for i in 1:(n - 1)
         xi = x[index]
@@ -171,7 +171,7 @@ function transform_with(flag::LogJacFlag, t::UnitVectorNorm, x::AbstractVector,
     (; n, chi_prior) = t
     T = _ensure_float(eltype(x))
     log_r = zero(T)
-    y = Vector{T}(undef, n)
+    y = similar(x, T, n)
     copyto!(y, 1, x, index, n)
     r = norm(y, 2)
     __normalize!(y, r)
@@ -230,7 +230,7 @@ function transform_with(flag::LogJacFlag, t::UnitSimplex, x::AbstractVector, ind
     T = _ensure_float(eltype(x))
     ℓ = logjac_zero(flag, T)
     stick = one(T)
-    y = Vector{T}(undef, n)
+    y = similar(x, T, n)
     @inbounds for i in 1:n-1
         xi = x[index]
         index += 1

src/utilities.jl

@@ -2,7 +2,23 @@
 ### logistic and logit
 ###
 
-function logistic_logjac(x::Real)
+
+"""
+    $(SIGNATURES)
+Method used in place of `Base.getindex` within this package.
+When `a <: Reactant.AnyTracedRArray` this adds a `@allowscalar` annotation so that the function can 
+be compiled with Reactant. When `a` is not a `Reactant.AnyTracedRArray` it simply returns `a[i]`.
+
+!!! warn
+    This is necessary because by default Reactant does not allow scalar indexing of arrays unless you
+    opt in. Note that often `Reactant` is able to raise the scalar indexing to the level of the whole
+    array so this operation is not necessarily slow, but there are no guarantees. 
+"""
+Base.@propagate_inbounds function tv_getindex(a, i)
+    return a[i]
+end
+
+function logistic_logjac(x::Number)
     mx = -abs(x)
     mx - 2*log1pexp(mx)
 end

test/Project.toml

@@ -12,6 +12,7 @@ LogDensityProblemsAD = "996a588d-648d-4e1f-a8f0-a84b347e47b1"
 OffsetArrays = "6fe1bfb0-de20-5000-8ca7-80f57d26f881"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+Reactant = "3c362404-f566-11ee-1572-e11a4b42c853"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 TransformVariables = "84d833dd-6860-57f9-a1a7-6da5db126cff"