Make logits shape more flexible in CategoricalHead.
Make logits shape more flexible in CategoricalHead.

PiperOrigin-RevId: 382584040
Change-Id: If57ede59bd7438829c00605a7a071e1be4683a9e

deepmind/acme,deepmind/acme