Migrate RLax squashed gaussian to use Distrax. Explicitly broadcast shapes in Distrax scalar affine to avoid rank promotion errors.
Migrate RLax squashed gaussian to use Distrax. Explicitly broadcast shapes in Distrax scalar affine to avoid rank promotion errors.

PiperOrigin-RevId: 441322360

deepmind/distrax,deepmind/distrax