Gumbel argmax for GPU sampling

MurrellGroup · Nov 27, 2024 · d8a397c · d8a397c
1 parent ee77d4b
commit d8a397c
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 0 deletions.
diff --git a/Project.toml b/Project.toml
@@ -5,10 +5,12 @@ version = "1.0.0-DEV"
 
 [deps]
 NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 
 [compat]
 NNlib = "0.9"
+Random = "1.11.0"
 StatsBase = "0.34"
 julia = "1.9"
 

diff --git a/src/samplers.jl b/src/samplers.jl
@@ -1,3 +1,10 @@
+#Gumbel argmax trick for GPU sampling
+function samplelogits(logits::AbstractVector)
+    u = similar(logits)
+    rand!(u)
+    return argmax(-log.(-log.(u)) .+ logits)
+end
+
 #To do: refactor into a combination of modified_softmax and sample. This way we can viz the result of the modified logits without having to sample.
 #This won't be visible to the user. Any method that doesn't fit this interface can be implemented directly.