don't vectorize all the logits processor state

Author: joelpaulkoch

lib/bumblebee/text/generation.ex

@@ -687,15 +687,15 @@ defmodule Bumblebee.Text.Generation do
       input_length: state.input_length
     }
 
-    {logits, new_logits_processor_state} =
-      logits_processor_process_fun.(logits, context, Nx.vectorize(state.logits_processor_state, :batch))
+    {logits, logits_processor_state} =
+      logits_processor_process_fun.(
+        logits,
+        context,
+        state.logits_processor_state
+      )
 
     logits = Nx.devectorize(logits, keep_names: false)
 
-    logits_processor_state =
-      Nx.devectorize(new_logits_processor_state, keep_names: false)
-
-
     {logits, %{state | logits_processor_state: logits_processor_state}}
   end
 

test/bumblebee/text/generation_test.exs

@@ -248,15 +248,15 @@ defmodule Bumblebee.Text.GenerationTest do
 
     @impl Bumblebee.LogitsProcessor
     def process(_logits_processor, state, logits, _context) do
-      sfp_state = state.sfp_state
-      logits = enforce_token(logits, sfp_state.next_enforced_token_id)
+      next_enforced_token_id = Nx.vectorize(state.sfp_state.next_enforced_token_id, :batch)
 
-      sfp_state = %{
-        sfp_state
-        | next_enforced_token_id: Nx.add(sfp_state.next_enforced_token_id, 1)
-      }
+      logits = enforce_token(logits, next_enforced_token_id)
+
+      next_enforced_token_id =
+        Nx.add(next_enforced_token_id, 1)
+        |> Nx.devectorize(keep_names: false)
 
-      state = %{state | sfp_state: sfp_state}
+      state = put_in(state.sfp_state.next_enforced_token_id, next_enforced_token_id)
 
       {logits, state}
     end