Static v0.8 compatibility

Author: oschulz

Project.toml

@@ -48,6 +48,6 @@ NaNMath = "0.3, 1"
 PrettyPrinting = "0.3, 0.4"
 Reexport = "1"
 SpecialFunctions = "2"
-Static = "0.5, 0.6"
+Static = "0.5, 0.6, 0.8"
 Tricks = "0.1"
 julia = "1.3"

src/MeasureBase.jl

@@ -32,7 +32,7 @@ using PrettyPrinting
 const Pretty = PrettyPrinting
 
 using ChainRulesCore
-using FillArrays
+import FillArrays
 using Static
 using FunctionChains
 
@@ -106,6 +106,7 @@ using Compat
 
 using IrrationalConstants
 
+include("static.jl")
 include("smf.jl")
 include("getdof.jl")
 include("transport.jl")

src/combinators/power.jl

@@ -1,58 +1,67 @@
 import Base
-using FillArrays: Fill
-# """
-# A power measure is a product of a measure with itself. The number of elements in
-# the product determines the dimensionality of the resulting support.
 
-# Note that power measures are only well-defined for integer powers.
+export PowerMeasure
 
-# The nth power of a measure μ can be written μ^x.
-# """
-# PowerMeasure{M,N,D} = ProductMeasure{Fill{M,N,D}}
+"""
+    struct PowerMeasure{M,...} <: AbstractProductMeasure
 
-export PowerMeasure
+    A power measure is a product of a measure with itself. The number of elements in
+    the product determines the dimensionality of the resulting support.
 
+    Note that power measures are only well-defined for integer powers.
+
+    The nth power of a measure μ can be written μ^x.
+"""
 struct PowerMeasure{M,A} <: AbstractProductMeasure
     parent::M
     axes::A
 end
 
+dslength(μ::PowerMeasure) = prod(dssize(μ))
+dssize(μ::PowerMeasure) = map(dslength, μ.axes)
+
 function Pretty.tile(μ::PowerMeasure)
     sz = length.(μ.axes)
     arg1 = Pretty.tile(μ.parent)
     arg2 = Pretty.tile(length(sz) == 1 ? only(sz) : sz)
     return Pretty.pair_layout(arg1, arg2; sep = " ^ ")
 end
 
+# ToDo: Make rand return static arrays for statically-sized power measures.
+
+_cartidxs(axs::Tuple{Vararg{<:AbstractUnitRange,N}}) where {N} = CartesianIndices(map(_dynamic, axs))
+
 function Base.rand(
     rng::AbstractRNG,
     ::Type{T},
     d::PowerMeasure{M},
 ) where {T,M<:AbstractMeasure}
-    map(CartesianIndices(d.axes)) do _
+    map(_cartidxs(d.axes)) do _
         rand(rng, T, d.parent)
     end
 end
 
 function Base.rand(rng::AbstractRNG, ::Type{T}, d::PowerMeasure) where {T}
-    map(CartesianIndices(d.axes)) do _
+    map(_cartidxs(d.axes)) do _
         rand(rng, d.parent)
     end
 end
 
+@inline _pm_axes(sz::Tuple{Vararg{<:IntegerLike,N}}) where N = map(one_to, sz)
+@inline _pm_axes(axs::Tuple{Vararg{<:AbstractUnitRange,N}}) where N = axs
+
 @inline function powermeasure(x::T, sz::Tuple{Vararg{<:Any,N}}) where {T,N}
-    a = axes(Fill{T,N}(x, sz))
-    A = typeof(a)
-    PowerMeasure{T,A}(x, a)
+    PowerMeasure(x, _pm_axes(sz))
 end
 
-marginals(d::PowerMeasure) = Fill(d.parent, d.axes)
+
+marginals(d::PowerMeasure) = fill_with(d.parent, d.axes)
 
 function Base.:^(μ::AbstractMeasure, dims::Tuple{Vararg{<:AbstractArray,N}}) where {N}
     powermeasure(μ, dims)
 end
 
-Base.:^(μ::AbstractMeasure, dims::Tuple) = powermeasure(μ, Base.OneTo.(dims))
+Base.:^(μ::AbstractMeasure, dims::Tuple) = powermeasure(μ, one_to.(dims))
 Base.:^(μ::AbstractMeasure, n) = powermeasure(μ, (n,))
 
 # Base.show(io::IO, d::PowerMeasure) = print(io, d.parent, " ^ ", size(d.xs))
@@ -76,7 +85,7 @@ end
 end
 
 @inline function logdensity_def(
-    d::PowerMeasure{M,Tuple{Base.OneTo{StaticInt{N}}}},
+    d::PowerMeasure{M,Tuple{StaticOneTo{N}}},
     x,
 ) where {M,N}
     parent = d.parent
@@ -86,7 +95,7 @@ end
 end
 
 @inline function logdensity_def(
-    d::PowerMeasure{M,NTuple{N,Base.OneTo{StaticInt{0}}}},
+    d::PowerMeasure{M,NTuple{N,StaticOneTo{0}}},
     x,
 ) where {M,N}
     static(0.0)
@@ -110,7 +119,7 @@ end
 
 @inline getdof(μ::PowerMeasure) = getdof(μ.parent) * prod(map(length, μ.axes))
 
-@inline function getdof(::PowerMeasure{<:Any,NTuple{N,Base.OneTo{StaticInt{0}}}}) where {N}
+@inline function getdof(::PowerMeasure{<:Any,NTuple{N,StaticOneTo{0}}}) where {N}
     static(0)
 end
 
@@ -135,7 +144,7 @@ logdensity_def(::PowerMeasure{P}, x) where {P<:PrimitiveMeasure} = static(0.0)
 
 # To avoid ambiguities
 function logdensity_def(
-    ::PowerMeasure{P,Tuple{Vararg{Base.OneTo{Static.StaticInt{0}},N}}},
+    ::PowerMeasure{P,Tuple{Vararg{StaticOneTo{0},N}}},
     x,
 ) where {P<:PrimitiveMeasure,N}
     static(0.0)

src/combinators/smart-constructors.jl

@@ -37,7 +37,7 @@ end
 ###############################################################################
 # ProductMeasure
 
-productmeasure(mar::Fill) = powermeasure(mar.value, mar.axes)
+productmeasure(mar::FillArrays.Fill) = powermeasure(mar.value, mar.axes)
 
 function productmeasure(mar::ReadonlyMappedArray{T,N,A,Returns{M}}) where {T,N,A,M}
     return powermeasure(mar.f.value, axes(mar.data))

src/domains.jl

@@ -116,7 +116,7 @@ struct Simplex <: CodimOne end
 
 function zeroset(::Simplex)
     f(x::AbstractArray{T}) where {T} = sum(x) - one(T)
-    ∇f(x::AbstractArray{T}) where {T} = Fill(one(T), size(x))
+    ∇f(x::AbstractArray{T}) where {T} = fill_with(one(T), size(x))
     ZeroSet(f, ∇f)
 end
 

src/standard/stdmeasure.jl

@@ -13,7 +13,7 @@ function transport_def(ν::StdMeasure, μ::PowerMeasure{<:StdMeasure}, x)
 end
 
 function transport_def(ν::PowerMeasure{<:StdMeasure}, μ::StdMeasure, x)
-    return Fill(transport_def(ν.parent, μ, only(x)), map(length, ν.axes)...)
+    return fill_with(transport_def(ν.parent, μ, only(x)), map(length, ν.axes))
 end
 
 function transport_def(
@@ -35,7 +35,7 @@ end
 # Implement transport_to(NU::Type{<:StdMeasure}, μ) and transport_to(ν, MU::Type{<:StdMeasure}):
 
 _std_measure(::Type{M}, ::StaticInt{1}) where {M<:StdMeasure} = M()
-_std_measure(::Type{M}, dof::Integer) where {M<:StdMeasure} = M()^dof
+_std_measure(::Type{M}, dof::IntegerLike) where {M<:StdMeasure} = M()^dof
 _std_measure_for(::Type{M}, μ::Any) where {M<:StdMeasure} = _std_measure(M, getdof(μ))
 
 function transport_to(::Type{NU}, μ) where {NU<:StdMeasure}
@@ -90,7 +90,7 @@ end
 @inline _offset_cumsum(s, x) = (s,)
 @inline _offset_cumsum(s) = ()
 
-function _stdvar_viewranges(μs::Tuple, startidx::Integer)
+function _stdvar_viewranges(μs::Tuple, startidx::IntegerLike)
     N = map(getdof, μs)
     offs = _offset_cumsum(startidx, N...)
     map((o, n) -> o:o+n-1, offs, N)

src/static.jl

@@ -0,0 +1,79 @@
+"""
+    MeasureBase.IntegerLike
+
+Equivalent to `Union{Integer,Static.StaticInt}`.
+"""
+const IntegerLike = Union{Integer,Static.StaticInt}
+
+
+"""
+    MeasureBase.one_to(n::IntegerLike)
+
+Creates a range from one to n.
+
+Returns an instance of `Base.OneTo` or `MeasureBase.StaticOneTo`, depending
+on the type of `n`.
+"""
+@inline one_to(n::Integer) = Base.OneTo(n)
+
+@static if isdefined(Static, :SOneTo)
+    @inline one_to(::Static.StaticInt{N}) where N = Static.SOneTo{N}()
+else
+    @inline one_to(::Static.StaticInt{N}) where N = Base.OneTo(StaticInt{N}())
+end
+
+
+"""
+    MeasureBase.StaticOneTo(n)
+
+A static range from one to N.
+
+StaticOneTo is a type alias for `Static.SOneTo{N}` (for Static.jl >= v0.8) or
+`Base.OneTo{Static.StaticInt{N}}` (older versions of Static.jl).
+"""
+@static if isdefined(Static, :SOneTo)
+    const StaticOneTo{N} = Static.SOneTo{N}
+else
+    const StaticOneTo{N} = Base.OneTo{Static.StaticInt{N}}
+end
+
+_dynamic(x::Number) = dynamic(x)
+_dynamic(::StaticOneTo{N}) where N = Base.OneTo(N)
+_dynamic(r::AbstractUnitRange) = minimum(r):maximum(r)
+
+"""
+    MeasureBase.fill_with(x, sz::NTuple{N,<:IntegerLike}) where N
+
+Creates an array of size `sz` filled with `x`.
+
+Returns an instance of `FillArrays.Fill`.
+"""
+function fill_with end
+
+@inline fill_with(x::T, sz::Tuple{Vararg{<:IntegerLike,N}}) where {T,N} = fill_with(x, map(one_to, sz))
+
+@inline function fill_with(x::T, axs::Tuple{Vararg{<:AbstractUnitRange,N}}) where {T,N}
+    # While `FillArrays.Fill` (mostly?) works with axes that are static unit
+    # ranges, some operations that automatic differentiation requires do fail
+    # on such instances of `Fill` (e.g. `reshape` from dynamic to static size).
+    # So need to use standard ranges for the axes for now:
+    dyn_axs = map(_dynamic, axs)
+    FillArrays.Fill(x, dyn_axs)
+end
+
+
+"""
+    MeasureBase.dslength(x)::IntegerLike
+
+Returns the length of `x` as a dynamic or static integer.
+"""
+dslength(x) = length(x)
+dslength(::Static.OptionallyStaticUnitRange{StaticInt{A}, StaticInt{B}}) where {A,B} = static(B - A + 1)
+
+
+"""
+    MeasureBase.dssize(x)::Tuple{Vararg{IntegerLike}}
+
+Returns the size of `x` as a tuple of dynamic or static integers.
+"""
+dssize(x) = size(x)

test/transport.jl

@@ -8,7 +8,7 @@ using LogExpFunctions: logit
 using ChainRulesTestUtils
 
 @testset "transport_to" begin
-    test_rrule(MeasureBase._origin_depth, pushfwd(exp, StdUniform()))
+    test_rrule(MeasureBase._origin_depth, pushfwd(exp, StdUniform()), output_tangent = static(0))
 
     for (f, μ) in [
         (logit, StdUniform())