[tx] Support top_k sampling

[agolajko]

Implements the top_k part of the sampling API requested in #533

Tests

1. test_top_k_filtering in test_generator.py: tests the core logic of the sampling
2. test_sample_top_k in test_api.py: checks the API can be called with the top_k parameter

Discussed with @pcmoritz on slack

[agolajko]

@gemini-code-assist review

[gemini-code-assist]

Code Review

This pull request successfully implements top_k sampling, including the core logic and API integration. The changes are well-tested. I've provided a few suggestions to enhance performance and improve the readability and robustness of the tests.

[gemini-code-assist]

Code Review

This pull request implements top_k sampling. The changes look good and include necessary logic in the generation pipeline and new tests. I've provided a few suggestions to improve code maintainability and performance. Specifically, I've suggested refactoring duplicated test code, simplifying assertions, and optimizing the apply_top_k function for better performance.

[gemini-code-assist]

The assertions to check the filtered logits can be simplified by comparing against an expected array. This makes the test more concise and easier to read.

Suggested change

	# Values below threshold should be -inf
	assert jnp.isinf(filtered[0]) and filtered[0] < 0
	assert jnp.isinf(filtered[1]) and filtered[1] < 0
	assert jnp.isinf(filtered[2]) and filtered[2] < 0
	# Top 2 values should be unchanged
	assert filtered[3] == 4.0
	assert filtered[4] == 5.0
	expected = jnp.array([-jnp.inf, -jnp.inf, -jnp.inf, 4.0, 5.0])
	assert jnp.array_equal(filtered, expected)

[tyler-griggs]

Ditto above

[agolajko]

@tyler-griggs lmk if you have any comments for top_k

Rebased to most recent main

[agolajko]

@pcmoritz @tyler-griggs lmk if you have thoughts, happy to rebase again to most recent before you merge

[tyler-griggs]

Thanks @agolajko !!

[tyler-griggs]

Ditto above

[pcmoritz]

/gemini review

[gemini-code-assist]

Code Review

This pull request successfully implements top_k sampling. The changes are well-structured, touching the API, type definitions, and the core generator logic. The implementation of apply_top_k_batch is efficient and JIT-friendly. The new tests in test_api.py and test_generator.py provide good coverage for the new functionality.

I have a few suggestions to improve the assertions in the tests to make them more robust and concise. Overall, this is a solid contribution.

[pcmoritz]

Thanks a lot for implementing this @agolajko , I implemented @tyler-griggs 's suggestion of a fast path if there is no top_k filtering, and also used jax.lax.top_k so we don't need to do the sorting :)

[pcmoritz]

I'll merge this now, but there are some improvements that could be made as a follow up:

• Bin max_top_k to avoid JIT recompilation (and we could also restrict top_k to be smaller than some reasonable value)
• Do sampling only among the top_k / max_top_k values, this will be more performant but requires reorganizing the code