Fix compression bug

joachimbrand · Oct 10, 2023 · 450a771 · 450a771
1 parent f83a7e6
commit 450a771
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/src/StochasticStyles/compression.jl b/src/StochasticStyles/compression.jl
@@ -10,13 +10,16 @@ struct ThresholdCompression{T} <: CompressionStrategy
 end
 ThresholdCompression() = ThresholdCompression(1)
 
+_set_value!(w, key, val) = w[key] = val
+_set_value!(w::Dict, key, val) = val ≠ 0 ? w[key] = val : delete!(w, key)
+
 function _threshold_compress!(t::ThresholdCompression, w, pairs)
     for (add, val) in pairs
         prob = abs(val) / t.threshold
         if prob < 1 # projection is only necessary if abs(val) < s.threshold
             val = ifelse(prob > rand(), t.threshold * sign(val), zero(val))
         end
-        w[add] = val
+        _set_value!(w, add, val)
     end
 end