fix issue #3: more robust indexing

all test should now pass

Author: simsurace

src/kernels.jl

@@ -32,14 +32,22 @@ end
 
 
 # BTRS algorithm, adapted from the tensorflow library (https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/kernels/random_binomial_op.cc)
-function kernel_BTRS!(A, count, prob, randstates)
+function kernel_BTRS!(A, count, prob, randstates, R1, R2, Rp, Ra, count_dim_larger_than_prob_dim)
     i = (blockIdx().x - 1) * blockDim().x + threadIdx().x
-    indices = CartesianIndices(A)
 
     @inbounds if i <= length(A)
-        I = indices[i].I
-        n = count[CartesianIndex(I[1:ndims(count)])]
-        p = prob[CartesianIndex(I[1:ndims(prob)])]
+        I  = Ra[i]
+        Ip = Rp[I[1]]
+        I1 = R1[Ip[1]]
+        I2 = R2[Ip[2]]
+
+        if count_dim_larger_than_prob_dim
+            n = count[CartesianIndex(I1, I2)]
+            p = prob[I1]
+        else
+            n = count[I1]
+            p = prob[CartesianIndex(I1, I2)]
+        end
 
         # wrong parameter values (currently disabled)
         # n < 0 && throw(ArgumentError("kernel_BTRS!: count must be a nonnegative integer."))

src/rand_binomial.jl

@@ -67,11 +67,25 @@ function rand_binom!(rng, A::BinomialArray, count::BinomialArray, prob::DenseCuA
         return A
     end
     if size(A)[1:ndims(count)] == size(count) && size(A)[1:ndims(prob)] == size(prob)
-        kernel  = @cuda name="BTRS_full" launch=false kernel_BTRS!(A, count, prob, rng.state)
+        count_dim_larger_than_prob_dim = ndims(count) > ndims(prob)
+        if count_dim_larger_than_prob_dim
+            R1 = CartesianIndices(prob) # indices for count
+            R2 = CartesianIndices(size(count)[ndims(prob)+1:end]) # indices for prob that are not included in R1
+            Rr = CartesianIndices(size(A)[ndims(count)+1:end]) # remaining indices in A
+        else
+            R1 = CartesianIndices(count) # indices for count
+            R2 = CartesianIndices(size(prob)[ndims(count)+1:end]) # indices for prob that are not included in R1
+            Rr = CartesianIndices(size(A)[ndims(prob)+1:end]) # remaining indices in A
+        end
+        Rp = CartesianIndices((length(R1), length(R2))) # indices for parameters
+        Ra = CartesianIndices((length(Rp), length(Rr))) # indices for parameters and A
+
+        kernel  = @cuda name="BTRS_full" launch=false kernel_BTRS!(A, count, prob, rng.state, R1, R2, Rp, Ra, count_dim_larger_than_prob_dim)
         config  = launch_configuration(kernel.fun)
         threads = Base.min(length(A), config.threads, 256) # strangely seems to be faster when defaulting to 256 threads
         blocks  = cld(length(A), threads)
-        kernel(A, count, prob, rng.state; threads=threads, blocks=blocks)
+
+        kernel(A, count, prob, rng.state, R1, R2, Rp, Ra, count_dim_larger_than_prob_dim; threads=threads, blocks=blocks)
     else
         throw(DimensionMismatch("`count` and `prob` need have size compatible with A"))
     end