Avoid rank promotion in stax biases (#4277)
Avoid rank promotion in stax biases (#4277)

* Avoid rank promotion in stax biases

* remove itertools
google/jax,google/jax,tensorflow/probability,google/jax,tensorflow/probability,google/jax