Chain of Responsibility for Logits Processing

[SamuelMarks]

Description

Sampling logic (sample_topk_topp_weighted) mixes distribution modification (temperature, filters) with random selection. Adapting from Hugging Face Transformer's design, this refactors inference_utils.py to use a LogitsProcessor chain (e.g., TemperatureLogitsWarper, TopKLogitsWarper, TopPLogitsWarper) and a separate Sampler. This helps not only readability now, but futureproofing of the introduction of many more [specialised] logit processors.

TL;DR this is a code smell:

  if algorithm == "greedy":
    return jnp.argmax(logits, axis=-1)
  elif algorithm == "weighted":
    return jax.random.categorical(rng, logits / temperature)
  elif algorithm == "nucleus":
    return sample_nucleus_topp_logits(logits, nucleus_topp, temperature, rng)
  elif algorithm == "topk":
    return sample_topk_logits(logits, topk, temperature, rng)
  elif algorithm == "composite":
    return sample_topk_topp_weighted(logits, topk, nucleus_topp, temperature, rng)
  else:
    raise ValueError(f"Sampling {algorithm=} not supported!")

Tests

CI

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!