Fixing up For and ProductMeasure

Author: cscherrer

src/combinators/for.jl

@@ -75,17 +75,13 @@ julia> For(eachrow(rand(4,2))) do x Normal(x[1], x[2]) end |> marginals |> colle
 ```
 
 """
-For(f, dims...) = ProductMeasure(f, dims)
+For(f, dims...) = ProductMeasure(i -> f(i...), zip(dims...))
 
 For(f, inds::AbstractArray) = ProductMeasure(f, inds)
 
-
+For(f, n::Int) = ProductMeasure(f, 1:n)
 For(f, dims::Int...) = ProductMeasure(i -> f(Tuple(i)...), CartesianIndices(dims))
 
-function For(f, inds::AbstractArray...)
-    g(x) = f(x...)
-    For(g, mappedarray(tuple, inds...))
-end
 
 function Base.eltype(d::ProductMeasure{F,I}) where {F,I<:AbstractArray}
     return eltype(d.f(first(d.pars)))

src/combinators/product.jl

@@ -134,13 +134,11 @@ function Base.rand(rng::AbstractRNG, ::Type{T}, d::ProductMeasure{F,I}) where {T
 end
 
 @propagate_inbounds function Random.rand!(rng::AbstractRNG, d::ProductMeasure, x::AbstractArray)
-    mar = marginals(d)
-    @boundscheck size(mar) == size(x) || throw(BoundsError)
-
-    @inbounds for j in eachindex(x)
-        x[j] = rand(rng, eltype(x), mar[j])
+    T = eltype(x)
+    for(j,m) in zip(eachindex(x), marginals(d))
+        @inbounds x[j] = rand(rng, T, m)
     end
-    x
+    return x
 end
 
 

test/runtests.jl

@@ -7,6 +7,8 @@ using DynamicIterators: trace, TimeLift
 using TransformVariables: transform, as𝕀, inverse
 
 using MeasureTheory
+using MeasureTheory: Const
+
 using Aqua
 Aqua.test_all(MeasureTheory; ambiguities=false, unbound_args=false)