Use functools.partial rather than JAX Partial.
Use functools.partial rather than JAX Partial.

Some very minor cleanup in topk too.

PiperOrigin-RevId: 303352055
Change-Id: I6cc1382b26f914577827ad116b445af9ffae2748

deepmind/dm-haiku,deepmind/dm-haiku