WIP: add state to logits processing

[joelpaulkoch]

this is the current state.

mix test test/bumblebee/text/generation_test.exs:135 fails with

  1) test with stateful logits processor with different batch sizes (Bumblebee.Text.GenerationTest)
     test/bumblebee/text/generation_test.exs:110
     ** (RuntimeError) unexpected vectorized axes in evaluator for operation :add: #Nx.Tensor<
       vectorized[batch: 1]
       s32[1]

       Nx.Defn.Expr
       tensor a        s32[1]
       b = reshape a   s32[1][1]
       c = reshape 1   s32[1][1]
       d = add b, c    s32[1][1]
     >
     code: %{token_ids: token_ids} = generate.(params, inputs)
     stacktrace:
       (nx 0.10.0) lib/nx/defn/evaluator.ex:435: Nx.Defn.Evaluator.eval_apply/4
       (nx 0.10.0) lib/nx/defn/evaluator.ex:274: Nx.Defn.Evaluator.eval/3
       (elixir 1.18.3) lib/enum.ex:1840: Enum."-map_reduce/3-lists^mapfoldl/2-0-"/3
       (elixir 1.18.3) lib/enum.ex:1840: Enum."-map_reduce/3-lists^mapfoldl/2-0-"/3
       (nx 0.10.0) lib/nx/container.ex:113: Nx.Container.Tuple.traverse/3
       (nx 0.10.0) lib/nx/defn/evaluator.ex:386: Nx.Defn.Evaluator.eval_apply/4
       (nx 0.10.0) lib/nx/defn/evaluator.ex:274: Nx.Defn.Evaluator.eval/3
       (nx 0.10.0) lib/nx/defn/evaluator.ex:339: Nx.Defn.Evaluator.eval_apply/4
       (nx 0.10.0) lib/nx/defn/evaluator.ex:274: Nx.Defn.Evaluator.eval/3
       (elixir 1.18.3) lib/enum.ex:1840: Enum."-map_reduce/3-lists^mapfoldl/2-0-"/3
       (nx 0.10.0) lib/nx/defn/evaluator.ex:439: Nx.Defn.Evaluator.eval_apply/4
       (nx 0.10.0) lib/nx/defn/evaluator.ex:274: Nx.Defn.Evaluator.eval/3
       (nx 0.10.0) lib/nx/container.ex:139: anonymous fn/3 in Nx.Container.Map.traverse/3
       (elixir 1.18.3) lib/enum.ex:1840: Enum."-map_reduce/3-lists^mapfoldl/2-0-"/3
       (nx 0.10.0) lib/nx/container.ex:138: Nx.Container.Map.traverse/3
       (nx 0.10.0) lib/nx/defn/evaluator.ex:82: anonymous fn/7 in Nx.Defn.Evaluator.__compile__/4
       (nx 0.10.0) lib/nx/defn/compiler.ex:134: Nx.Defn.Compiler.__jit__/4
       (nx 0.10.0) lib/nx/defn.ex:452: Nx.Defn.do_jit_apply/3
       test/bumblebee/text/generation_test.exs:159: (test)

However, this seems to be Nx.Defn.Evaluator specific (which is used in the tests). When switching to EXLA as compiler the test passes.

[joelpaulkoch]

Nx.Defn.Evaluator doesn't like adding on vectorized next_enforced_token_id

[joelpaulkoch]

the test passes when using EXLA as compiler

Suggested change

	%{token_ids: token_ids} = generate.(params, inputs)
	%{token_ids: token_ids} = Nx.Defn.jit_apply(generate, [params, inputs], compiler: EXLA)

[joelpaulkoch]

the test passes when using EXLA as compiler

Suggested change

	%{token_ids: token_ids} = generate.(params, inputs)
	%{token_ids: token_ids} = Nx.Defn.jit_apply(generate, [params, inputs], compiler: EXLA)