Implement top_k sampling

[AeroScripts]

This seems to work quite well in combination with top_p, I posted some comparisons on discord.

Using "hk.with_rng" to keep the seed consistent across runs makes it much easier to compare changes in the sampling methods, I used this to compare with and without top_k:

Prompt

In a shocking finding, scientists discovered a herd of unicorns living in a remote, previously unexplored valley, in the Andes Mountains. Even more surprising to the researchers was the fact that the unicorns spoke perfect English. 

top_p=0.9, seed=0

They came from a forgotten time period, pre-dating even the Jurassic period, and carried themselves as if they belonged in the upper echelons of society. The scientists, Pym Verbrugge and Daisy Kowalewski of Edinburgh University, decided to name the new breed 'Ungulata', meaning '1. the class of hoofed mammals with one or more extra toes; 2. a type of horse'.

Ungulata lived in a valley near Bolivia, which was itself located in South America. The discovery came as a total surprise to the scientists, who were out exploring the region for other (strictly scientific) reasons. They found the valley locked off by a barrier of snow, and upon crossing it, were overcome by an 'intensely fragrant atmosphere' before entering the valley.

On reaching the valley floor, the scientists noticed tracks in the snow leading up to a large group of snow covered boulders. The tracks appeared to have been made by a group of four animals; however, upon closer inspection, they realised that the tracks were actually of six. The animals who had been using the tracks - clearly evident at the end of each track - had been on their hind legs.

top_p=0.9+top_k=50, seed=0

They came from a land called Pangaea and were traveling through the mountains to settle down in the United States, Europe and Australia.

Unicorns are mythical, but there is some evidence that they exist. The evidence has long been held to be too speculative to prove they exist, but as global warming and deforestation continue to threaten the habitats of the species, the evidence is growing.

The discovery of unicorns in Peru was made in August 2012, but was first reported in January. Although their existence is not yet officially confirmed, the species is known to scientists as Boophone and is scientifically called B. unicornis. (The name “unicorn” is misleading because it applies to a different creature—the mammal genus Equus, specifically a horse.)

top_p=0.9, seed=1

How did this intelligent species of snow white animal with large, incisors evolve in such a small and isolated environment?

The now extinct mammals were first discovered in 1989 during an expedition to Peru. However, the discovery of the new species of unicorns took the world by surprise when, at the end of the year, scientists and others confirmed that a herd of the animals, found living on the shores of Lake Titicaca, were actually alive. The herd included a mother and her newborn calf.

#1 – Werewolf

They were beautiful, exotic, intelligent, carnivorous, and savage. Yes, the world’s most famous werewolves are the result of manmade genetics,

top_p=0.9+top_k=50, seed=1

How did this intelligent species of animals come to live in such an isolated area?

When the scientists arrived, they were greeted with an open invitation to come inside and meet with their hosts, the unicorns. After several days of trekking in the Andes Mountains, the team finally got the chance to talk with the creatures. But the animals were very shy and wouldn’t come anywhere close to the team’s tents.

As the team explored the valley further, they were startled to discover a herd of the unicorns living there. It was a very long trip back from their base camp when the team finally came across their animals. Many of the unicorns were walking around in circles, clearly lost, and some were sitting down next to the team’s tent, as if in a lost puppy moment.

It seems to stay on topic better with different prompts so I've enabled it by default.

top_k can be set to None to disable. I wanted to implement this using slices, but couldn't get it working properly. jax.lax.dynamic_slice isn't playing nicely. However, jnp.where produces identical results

[kingoflolz]

Looks good, can you just confirm that it doesn't break if you set top_k=None in infer?

[AeroScripts]

Yes I tested that and it works fine 👍

[kingoflolz]

Thanks!