Bug in predict_class cpu version

[khosravipasha]

For Mnist-1M.circuit from zoo, seems the cpu version returning wrong labels. For other zoo circuits like mnist.circuit or mnist-meduim.circuit they seem to fully match.

Code to reproduce:

using Revise
using LogicCircuits
using ProbabilisticCircuits
using DiscriminativeCircuits
using MLDatasets
using DataFrames: DataFrame
using Statistics

function mnist()
    # transposing makes slicing by variable much much faster
    # need to take a copy to physically move the data around
    train_x = collect(Float32, transpose(reshape(MNIST.traintensor(), 28*28, :)))
    test_x  = collect(Float32, transpose(reshape(MNIST.testtensor(), 28*28, :)))
    
    train_x = DataFrame(train_x, :auto)
    test_x = DataFrame(test_x, :auto)
    
    train_y::Vector{Int32} = MNIST.trainlabels()
    test_y::Vector{Int32}  = MNIST.testlabels()
    
    return train_x, train_y, test_x, test_y
end


function main()
    circuit_name = "mnist-1M.circuit"
    lc = zoo_dc(circuit_name)
    println("Loaded zoo($(circuit_name) with $(num_classes(lc)) classes")
    @time begin
        mnist_train, train_y, mnist_test, test_y = mnist();
        mnist_train, _, mnist_test = threshold(mnist_train[:, 1:end-1], nothing, mnist_test[:, 1:end-1], 0.5);
    end

    # WARMup
    predict_class(lc, num_classes(lc), mnist_train[1:3,:])
    predict_class(lc, num_classes(lc), to_gpu(mnist_train[1:3,:]))

    @time pred = predict_class(lc, num_classes(lc), mnist_train)
    pred .-= 1 # does not return lables its return the index of the label
    acc = mean( pred .== train_y )
    println("\t [CPU] Accuray = $( acc )")

    @time pred_gpu = predict_class(lc, num_classes(lc), to_gpu(mnist_train))
    pred_gpu .-=1
    acc = mean( to_cpu(pred_gpu) .== train_y )
    println("\t [GPU] Accuray = $( acc )")

    println("CPU - GPU Match $( mean( to_cpu(pred_gpu) .== pred ) );")

    println(train_y[1:10])
    println(pred[1:10])
    println(to_cpu(pred_gpu)[1:10])
end

output:

Loaded zoo(mnist-1M.circuit with 10 classes
  1.617942 seconds (27.09 k allocations: 921.044 MiB, 14.60% gc time)
  2.720844 seconds (1.22 M allocations: 899.938 MiB, 11.95% gc time)
         [CPU] Accuray = 0.061316666666666665
  0.210986 seconds (1.27 M allocations: 119.404 MiB)
         [GPU] Accuray = 0.8428
CPU - GPU Match 0.08871666666666667;
Int32[5, 0, 4, 1, 9, 2, 1, 3, 1, 4]
[0, 1, 4, 0, 0, 0, 0, 0, 0, 0]
[3, 0, 4, 1, 9, 2, 1, 3, 1, 4]