fix unweighted combination

gbaydin · gbaydin · commit 3edbce68d372 · 2022-02-20T19:22:30.000Z
diff --git a/src/DiffSharp.Core/Distributions.fs b/src/DiffSharp.Core/Distributions.fs
@@ -225,7 +225,7 @@ type Categorical(?probs:Tensor, ?logits:Tensor) =
 
 
 /// <summary>Represents an Empirical distribution.</summary>
-type Empirical<'T when 'T:equality>(values:seq<'T>, ?weights:Tensor, ?logWeights:Tensor, ?combineDuplicates:bool) =
+type Empirical<'T when 'T:equality>(values:seq<'T>, ?weights:Tensor, ?logWeights:Tensor, ?combineDuplicates:bool, ?device:Device, ?dtype:Dtype, ?backend:Backend) =
     inherit Distribution<'T>()
     let _categorical, _weighted =
         match weights, logWeights with
@@ -239,7 +239,7 @@ type Empirical<'T when 'T:equality>(values:seq<'T>, ?weights:Tensor, ?logWeights
     let _valuesTensor =
             lazy(try _values |> Array.map (fun v -> box v :?> Tensor) |> dsharp.stack
                     with | _ -> 
-                        try _values |> Array.map (dsharp.tensor) |> dsharp.stack
+                        try _values |> Array.map (dsharp.tensor(device=defaultArg device Device.Default, backend=defaultArg backend Backend.Default, dtype=defaultArg dtype Dtype.Default)) |> dsharp.stack
                         with | _ -> failwith "Not supported because Empirical does not hold values that are Tensors or can be converted to Tensors")
     do
         let combineDuplicates = defaultArg combineDuplicates false
@@ -256,7 +256,8 @@ type Empirical<'T when 'T:equality>(values:seq<'T>, ?weights:Tensor, ?logWeights
                     Dictionary.copyKeys uniques, dsharp.stack(Dictionary.copyValues uniques).view(-1)
                 else
                     let vals, counts = _values |> Array.getUniqueCounts false
-                    vals, dsharp.tensor(counts)
+                    let c = dsharp.tensor(counts, device=defaultArg device Device.Default, backend=defaultArg backend Backend.Default, dtype=defaultArg dtype Dtype.Default)
+                    vals, probsToLogits (c/c.sum()) false
             _values <- newValues
             _categorical <- Categorical(logits=newLogWeights)
         _weighted <- not (Seq.allEqual (_categorical.probs.unstack()))
diff --git a/tests/DiffSharp.Tests/TestDistributions.fs b/tests/DiffSharp.Tests/TestDistributions.fs
@@ -369,6 +369,23 @@ type TestDistributions () =
             Assert.CheckEqual(distUnweightedCombinedModeCorrect, distUnweightedCombinedMode)
             Assert.CheckEqual(distUnweightedCombinedLengthCorrect, distUnweightedCombinedLength)
 
+    [<Test>]
+    member _.TestDistributionsEmpiricalCombineDuplicatesUnweighted () =
+        for combo in Combos.AllDevicesAndBackendsFloat32 do
+            let values = combo.tensor([0,1,1,1,2,2])
+
+            let dist = Empirical(values.unstack(), combineDuplicates=true, device=combo.device, backend=combo.backend, dtype=combo.dtype)
+            let values = dist.valuesTensor
+            let weights = dist.weights
+            let valuesCorrect = combo.tensor([0,1,2])
+            let weightsCorrect = combo.tensor([1./6., 3./6., 2./6.])
+            printfn "%A" values.dtype
+            printfn "%A" weights.dtype
+            printfn "%A" valuesCorrect.dtype
+            printfn "%A" weightsCorrect.dtype
+            Assert.True(valuesCorrect.allclose(values, 0.1))
+            Assert.True(weightsCorrect.allclose(weights, 0.1))
+
     [<Test>]
     member _.TestDistributionsEmpiricalResampleFilter () =
         for combo in Combos.AllDevicesAndBackendsFloat32 do
