Segmented_sum in Knet

Specific Domains Machine Learning
,
1

Hello, I have started to play with Knet, but I am unable to reproduce my work related to multi-instance learning in it, since I need segmented_sum operation from tensorflow. I have created a small example

using Knet
using MLLabelUtils

function segmented_sum(x,bags)
xx=zeros(length(bags),size(x,2))
for (i,bag) in enumerate(bags)
xx[i,:]=sum(x[bag,:],1)
end
return(xx)
end

predict(w,x,bags) = segmented_sum(xw[1] .+ w[2],bags)w[3] .+ w[4]
function softmaxloss(w,x,bags,ygold)
ypred = predict(w,x,bags)
ynorm = ypred .- logsumexp(ypred,2)
-sum(ygold .* ynorm) / size(ygold,1)
end
lossgradient = grad(softmaxloss)
x=randn(100,5)
bags=[i:i+9 for i in 1:10:100]
ygold=convertlabel(LabelEnc.OneOfK{Float64},rand(1:2,10), obsdim = 1)
d=size(x,2)
k=4
w=[randn(d,k),randn(1,k),randn(k,2),randn(1,2)]
g = lossgradient(w,x,bags,ygold)

The problematic part is xx[i,:]=sum(x[bag,:],1) where I get an error
Cannot convert an object of type AutoGrad.Rec{Array{Float64,2}} to an object of type Float64
Can anyone point me, how to fix this?

I have found that I can implement segmented_sum as

function segmented_sum(x,bags)
mapreduce(b->sum(x[b,:],1),vcat,bags)
end

But this is terribly slow.

Thank for any help.