Update DiscreteValuedDistribution to track recent upstream changes to TFP.
Update DiscreteValuedDistribution to track recent upstream changes to TFP.

TFP now derives batch shapes automatically from parameter annotations. As part of this change, Categorical (and other distributions) calls  `self._batch_shape_tensor` using named (rather than positional) args:
https://github.com/tensorflow/probability/commit/390378db8ca9cfeee5ff5b352f741770d03f90d0#diff-0581143214632da986e20885048206f9bb15f9199034b5c7b3cb59d35401466aL260

The current implementation of `DiscreteValuedDistribution._batch_shape_tensor` doesn't expect `logits` as a keyword arg, thus `DiscreteValuedDistribution._sample_n` is broken with recent TFP snapshots. This change updates `DiscreteValuedDistribution` to rely on the now-generic Distribution implementation of `_batch_shape_tensor`.

PiperOrigin-RevId: 380587296
Change-Id: I58bdd31b08947ad82be0e06e6e67821a232afa20

deepmind/acme,deepmind/acme