Method generate

src/MaxText/maxengine.py:915–951

Public API for generate that updates page state outside JIT.

(
      self,
      params: Params,
      decode_state: DecodeState,
      sampler: Callable[[Any], Any] | None = None,  # pylint: disable=unused-argument
      rng: PRNGKeyType | None = None,
      algorithm: str | None = None,
      topk: int | None = None,
      nucleus_topp: float | None = None,
      temperature: float | None = None,
  )

  # Public non-JIT generate method that updates page state
  def generate(
      self,
      params: Params,
      decode_state: DecodeState,
      sampler: Callable[[Any], Any] | None = None,  # pylint: disable=unused-argument
      rng: PRNGKeyType | None = None,
      algorithm: str | None = None,
      topk: int | None = None,
      nucleus_topp: float | None = None,
      temperature: float | None = None,
  ) -> tuple[DecodeState, engine_api.ResultTokens]:
    """Public API for generate that updates page state outside JIT."""

    # Update page state before JIT call
    if self.page_manager is not None and self.page_state is not None:
      self.page_state = self.page_manager.update_decode_pages(self.page_state)

    # Sample rng before JIT call
    if rng is None:
      if self.rng is None:
        self.rng = jax.random.PRNGKey(0)
      self.rng, rng = jax.random.split(self.rng)

    # Call JIT-compiled version with current state
    new_state, result = self._generate_jit(
        params=params,
        decode_state=decode_state,
        sampler=sampler,
        page_state=self.page_state,
        rng=rng,
        algorithm=algorithm,
        topk=topk,
        nucleus_topp=nucleus_topp,
        temperature=temperature,
    )

    return max_utils.unbox_logicallypartioned(new_state), result

  @functools.partial(
      jax.jit, static_argnums=(0,), donate_argnums=(2,), static_argnames=("algorithm", "topk", "nucleus_topp")