elixir-nx · jonatanklosko · May 4, 2023 · Apr 21, 2023 · Apr 28, 2023 · May 4, 2023
diff --git a/lib/bumblebee.ex b/lib/bumblebee.ex
@@ -117,6 +117,10 @@ defmodule Bumblebee do
     "GPT2ForTokenClassification" => {Bumblebee.Text.Gpt2, :for_token_classification},
     "GPT2LMHeadModel" => {Bumblebee.Text.Gpt2, :for_causal_language_modeling},
     "GPT2Model" => {BumbleBee.Text.Gpt2, :base},
+    "GPTNeoXModel" => {Bumblebee.Text.GptNeoX, :base},
+    "GPTNeoXForCausalLM" => {Bumblebee.Text.GptNeoX, :for_causal_language_modeling},
+    "GPTNeoXForSequenceClassification" => {Bumblebee.Text.GptNeoX, :for_sequence_classification},
+    "GPTNeoXForTokenClassification" => {Bumblebee.Text.GptNeoX, :for_token_classification},
     "LayoutLMForMaskedLanguageModeling" =>
       {Bumblebee.Multimodal.LayoutLm, :for_masked_language_modeling},
     "LayoutLMForQuestionAnswering" => {Bumblebee.Multimodal.LayoutLm, :for_question_answering},
@@ -195,6 +199,7 @@ defmodule Bumblebee do
     "distilbert" => Bumblebee.Text.DistilbertTokenizer,
     "camembert" => Bumblebee.Text.CamembertTokenizer,
     "clip" => Bumblebee.Text.ClipTokenizer,
+    "gpt_neox" => Bumblebee.Text.GptNeoXTokenizer,
     "gpt2" => Bumblebee.Text.Gpt2Tokenizer,
     "layoutlm" => Bumblebee.Text.LayoutLmTokenizer,
     "llama" => Bumblebee.Text.LlamaTokenizer,

diff --git a/lib/bumblebee/audio/whisper.ex b/lib/bumblebee/audio/whisper.ex
@@ -428,11 +428,11 @@ defmodule Bumblebee.Audio.Whisper do
         layer_norm: [
           epsilon: 1.0e-5
         ],
-        norm_placement: :first,
         ffn: [
           intermediate_size: spec.encoder_intermediate_size,
           activation: spec.activation
         ],
+        block_type: :norm_first,
         output_hidden_states: spec.output_hidden_states,
         output_attentions: spec.output_attentions,
         name: join(name, "blocks")
@@ -477,11 +477,11 @@ defmodule Bumblebee.Audio.Whisper do
         layer_norm: [
           epsilon: 1.0e-5
         ],
-        norm_placement: :first,
         ffn: [
           intermediate_size: spec.decoder_intermediate_size,
           activation: spec.activation
         ],
+        block_type: :norm_first,
         output_hidden_states: spec.output_hidden_states,
         output_attentions: spec.output_attentions,
         name: join(name, "blocks")

diff --git a/lib/bumblebee/diffusion/layers/unet.ex b/lib/bumblebee/diffusion/layers/unet.ex
@@ -323,8 +323,8 @@ defmodule Bumblebee.Diffusion.Layers.UNet do
         epsilon: 1.0e-5
       ],
       dropout_rate: dropout,
-      norm_placement: :first,
       ffn: &ffn_geglu(&1, hidden_size, dropout: dropout, name: &2),
+      block_type: :norm_first,
       name: join(name, "blocks")
     )
     |> then(fn %{hidden_state: hidden_state} ->

diff --git a/lib/bumblebee/layers.ex b/lib/bumblebee/layers.ex
@@ -979,7 +979,7 @@ defmodule Bumblebee.Layers do
   end
 
   deftransformp create_sinusoidal_positions(max_positions, size, base) do
-    range = Nx.multiply(Nx.iota({div(size, 2)}), 2)
+    range = Nx.iota({div(size, 2)}) |> Nx.multiply(2) |> Nx.divide(size)
     inv_frequency = Nx.divide(1.0, Nx.pow(base, range))
 
     position = Nx.iota({max_positions})

diff --git a/lib/bumblebee/layers/decoder.ex b/lib/bumblebee/layers/decoder.ex
@@ -59,14 +59,14 @@ defmodule Bumblebee.Layers.Decoder do
     decoder_head_size =
       opts[:attention_head_size] || div(hidden_size, decoder_num_attention_heads)
 
-    encoder_head_size =
-      opts[:attention_head_size] || div(hidden_size, encoder_num_attention_heads)
-
     self_attention =
       attention_cache(batch_size, max_length, decoder_num_attention_heads, decoder_head_size)
 
     cross_attention =
       if encoder_sequence_length do
+        encoder_head_size =
+          opts[:attention_head_size] || div(hidden_size, encoder_num_attention_heads)
+
         attention_cache(
           batch_size,
           encoder_sequence_length,