feat: more progress

Author: avik-pal

examples/NanoGPT/Project.toml

@@ -14,7 +14,6 @@ Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Reactant = "3c362404-f566-11ee-1572-e11a4b42c853"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
-StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 
 [compat]
 Comonicon = "1"
@@ -30,4 +29,3 @@ Printf = "1.10"
 Random = "1.10"
 Reactant = "0.2.5"
 Statistics = "1.10"
-StatsBase = "0.34.3"

examples/NanoGPT/main.jl

@@ -1,3 +1,5 @@
+ENV["XLA_REACTANT_GPU_MEM_FRACTION"] = "0.98"
+
 using ConcreteStructs,
     MLUtils,
     Lux,
@@ -120,25 +122,35 @@ function (model::GPT2)(x, ps, st)
     return outputs, (; tok_emb=st_tok_emb, pos_emb=st_pos_emb, gpt_blocks=st_gpt_blocks)
 end
 
+#=
 dev = reactant_device(; force=true)
 rng = Random.default_rng()
 
 model = GPT2(;
     n_vocab=50304,
-    embed_dim=768,
+    embed_dim=1024,
     hidden_dim=3072,
     block_size=1024,
     n_layers=3,
     dropout_rate=0.0,
-    num_heads=12,
+    num_heads=16,
 )
 ps, st = Lux.setup(rng, model) |> dev;
 
-x = rand(1:50304, 1024, 32) |> dev;
+x = rand(1:50304, 48, 32) |> dev;
 
 @code_hlo model(x, ps, st)
 
+sumabs2first(layer, x, ps, st) = sum(abs2, first(layer(x, ps, st)))
+
+@code_hlo Enzyme.gradient(Reverse, sumabs2first, Const(model), x, ps, Const(st))
+=#
+
 # Use the model to generate some text.
+# function weighted_sample(items::AbstractVector, weights::AbstractVector)
+
+# end
+
 function generate_text(model, ps, st, seed; alphabet, output_length, sequence_length)
     dev = get_device((ps, st))
     @assert !(dev isa ReactantDevice) "Currently we don't support running inference of \

src/layers/embedding.jl

@@ -63,7 +63,7 @@ function (e::Embedding)(x::Union{Number,AbstractVector}, ps, st::NamedTuple)
 end
 function (e::Embedding)(x::AbstractArray, ps, st::NamedTuple)
     @argcheck Utils.eltype(x) <: Integer
-    y, stₙ = e(vec(x), ps, st)
+    y, stₙ = e(Utils.vec(x), ps, st)
     return reshape(y, :, size(x)...), stₙ
 end
 function (e::Embedding)(x::NTuple{N,T}, ps, st::NamedTuple) where {N,T}